"""Message types for the Swarm framework."""

from typing import TYPE_CHECKING, Any, Dict, List, Optional
from uuid import uuid4

from openai.types.chat import ChatCompletionMessageParam
from pydantic import BaseModel, field_validator

from swarm.chunk import TokenUsage
from swarm.providers.types import ToolCallDelta
from swarm.types import ContentType, MessageRole

if TYPE_CHECKING:
    from swarm.tool import Result


class BaseMessage(BaseModel):
    """Base message type with common fields"""

    id: Optional[str] = None
    content: List[Dict[str, Any]]
    sender: Optional[str] = None
    role: MessageRole
    content_type: ContentType = ContentType.TEXT
    function_call: Optional[Dict[str, Any]] = None
    component_type: Optional[Any] = None
    tool_calls: Optional[List[ToolCallDelta]] = None
    is_error: bool = False
    metadata: Optional[Dict[str, Any]] = (
        None  # Optional metadata for synthetic messages
    )

    def model_dump(self, *args: Any, **kwargs: Any) -> Dict[str, Any]:
        data = super().model_dump(*args, **kwargs)
        if "id" in data and hasattr(data["id"], "__class__"):
            data["id"] = str(data["id"])
        return data

    def to_openai_dict(self) -> Dict[str, Any]:
        """Convert message to OpenAI-compatible dict format"""
        data = self.model_dump(mode="json")
        data["id"] = str(data["id"])
        return {k: v for k, v in data.items() if v is not None}  # Remove None values


class SwarmAssistantMessage(BaseMessage):
    """Messages from the designer"""

    role: MessageRole = MessageRole.ASSISTANT
    usage: Optional[TokenUsage] = None

    @field_validator("content", mode="before")
    @classmethod
    def normalize_content(cls, v: Any) -> List[Dict[str, Any]]:
        """Normalize content to expected list format.

        Handles the case where content comes as a string from OpenAI API chunks
        and converts it to the expected list format.
        """
        if isinstance(v, str):
            return [{"type": "text", "text": v}] if v else []
        elif v is None:
            return []
        elif isinstance(v, list):
            return v
        else:
            # Fallback for unexpected types - convert to string representation
            return [{"type": "text", "text": str(v)}]

    def to_openai_dict(self) -> Dict[str, Any]:
        data = super().to_openai_dict()
        # Usage must never be sent back to the model
        data.pop("usage", None)
        if self.tool_calls:
            data["tool_calls"] = [
                {
                    "id": tc.id,
                    "type": tc.type,
                    "function": {
                        "name": tc.function.name,
                        "arguments": tc.function.arguments,
                    },
                    "index": tc.index,
                }
                for tc in self.tool_calls
            ]
        return data


class SwarmToolMessage(BaseMessage):
    """Messages from tool executions"""

    role: MessageRole = MessageRole.TOOL
    tool_call_id: str  # Required for tool messages
    tool_name: str  # Required for tool messages

    @field_validator("content", mode="before")
    @classmethod
    def normalize_content(cls, v: Any) -> List[Dict[str, Any]]:
        """Normalize content to expected list format.

        Handles the case where content comes as a string (e.g., JSON from tool results)
        and converts it to the expected list format.
        """
        if isinstance(v, str):
            return [{"type": "text", "text": v}] if v else []
        elif v is None:
            return []
        elif isinstance(v, list):
            return v
        else:
            # Fallback for unexpected types - convert to string representation
            return [{"type": "text", "text": str(v)}]

    @classmethod
    def from_result(
        cls, result: "Result", tool_call_id: str, tool_name: str
    ) -> "SwarmToolMessage":
        """Create SwarmToolMessage from a Result object"""
        return cls(
            id=uuid4(),
            content=[{"type": "text", "text": result.value}],
            tool_call_id=tool_call_id,
            tool_name=tool_name,
            component_type=result.component_type,
            content_type=result.content_type,
            is_error=result.is_error,
        )


class SwarmUserMessage(BaseMessage):
    """Messages from the user"""

    role: MessageRole = MessageRole.USER


SwarmMessage = SwarmAssistantMessage | SwarmToolMessage | SwarmUserMessage


def to_swarm_message(data: ChatCompletionMessageParam) -> SwarmMessage:
    """Create SwarmMessage from OpenAI format dict.

    Args:
        data: OpenAI format message dict

    Returns:
        SwarmMessage: Either SwarmAssistantMessage, SwarmToolMessage, or SwarmUserMessage

    Raises:
        ValueError: If role is not supported
    """

    # Helper function to convert content to dict format
    def normalize_content(content: Any) -> List[Dict[str, Any]]:
        if isinstance(content, str):
            # Convert string to dict format
            return [{"type": "text", "text": content}] if content else []
        elif isinstance(content, list):
            # Already in list format
            return content
        else:
            # Empty or None
            return []

    # Cast to Any to work with TypedDict union
    role = data.get("role", "")
    if role == "assistant":
        tool_calls = data.get("tool_calls")
        return SwarmAssistantMessage(
            id=str(uuid4()),
            content=normalize_content(data.get("content")),
            role=MessageRole.ASSISTANT,
            tool_calls=(
                [ToolCallDelta(**tc) for tc in tool_calls if isinstance(tc, dict)]
                if isinstance(tool_calls, list)
                else None
            ),
            sender=data.get("sender"),
        )
    elif role == "tool":
        return SwarmToolMessage(
            id=str(uuid4()),
            content=normalize_content(data.get("content")),
            role=MessageRole.TOOL,
            tool_call_id=data.get("tool_call_id", ""),
            tool_name=data.get("tool_name") or "unknown",
            sender=data.get("sender"),
        )
    elif role == "user":
        return SwarmUserMessage(
            id=str(uuid4()),
            content=normalize_content(data.get("content")),
            role=MessageRole.USER,
            sender=data.get("sender"),
        )
    else:
        raise ValueError(f"Unsupported message role: {role}")
