"""Streaming chunk types for the Swarm framework."""

from typing import Any, Dict, List, Literal, Optional, Union

from openai.types.chat.chat_completion_message_tool_call import (
    ChatCompletionMessageToolCall,
)
from pydantic import BaseModel, Field

from swarm.guardrails import GuardrailResult
from swarm.providers.types import ToolCallDelta
from swarm.types import ContentType, MessageRole


class BaseChunk(BaseModel):
    """Base model for all message chunks"""

    id: str
    content: Optional[str] = None
    function_call: Optional[Dict[str, Any]] = None  # Deprecated
    refusal: Optional[Any] = None
    tool_calls: Optional[List[ToolCallDelta]] = None
    sender: Optional[str] = None
    parent_id: Optional[str] = None


class ToolCallChunk(BaseChunk):
    """Tool call chunks - either complete or streaming"""

    role: MessageRole = MessageRole.ASSISTANT
    tool_calls: List[ToolCallDelta] = Field(
        default_factory=List[ToolCallDelta]
    )  # Override to make required with default


class AssistantMessageChunk(BaseChunk):
    """Assistant message chunks"""

    role: MessageRole = Field(default=MessageRole.ASSISTANT)


class ToolResponseContent(BaseModel):
    """Individual tool response content"""

    id: str
    role: MessageRole = MessageRole.TOOL
    tool_call_id: str
    tool_name: str
    content: str
    component_type: Optional[Any] = None
    content_type: ContentType = ContentType.TEXT
    is_error: bool = False
    metadata: Optional[Dict[str, Any]] = None


def _empty_tool_response_messages() -> (
    Dict[Literal["messages"], List[ToolResponseContent]]
):
    return {"messages": []}


class ToolResponseChunk(BaseModel):
    """Tool response chunks"""

    tool_response: Dict[Literal["messages"], List[ToolResponseContent]] = Field(
        default_factory=_empty_tool_response_messages
    )


class TokenUsage(BaseModel):
    """Typed token usage record for billing/cost tracking."""

    model: str
    input_tokens: int = 0
    output_tokens: int = 0


class ExecutionMetadata(BaseModel):
    """Execution metadata for autonomous loop observability."""

    agent_transitions: int = 0
    tool_executions: int = 0
    thinking_turns: int = 0
    agent_completed: bool = False
    agent_paused: bool = False
    completion_reason: Optional[str] = None
    accumulated_usage: List[TokenUsage] = Field(default_factory=list)
    output_guardrail_results: List[GuardrailResult] = Field(default_factory=list)
    dangling_tool_calls: List[ChatCompletionMessageToolCall] = Field(
        default_factory=list
    )


class DelimiterChunk(BaseModel):
    """Delimiter chunks"""

    delim: Literal["start", "end", "complete", "stopped"]
    execution_metadata: Optional["ExecutionMetadata"] = None


# Union type for all possible chunks
ChunkType = Union[
    AssistantMessageChunk,
    ToolCallChunk,
    ToolResponseChunk,
    DelimiterChunk,
]
