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

from pydantic import BaseModel, ConfigDict, field_validator, model_validator

SYSTEM = "system"
USER = "user"
ASSISTANT = "assistant"
TOOL = "tool"


class ContentItem(BaseModel):
    text: Optional[str] = None
    image: Optional[str] = None
    file: Optional[str] = None
    audio: Optional[Union[str, dict]] = None
    video: Optional[Union[str, list]] = None

    @model_validator(mode="after")
    def check_exclusivity(self):
        provided_fields = 0
        if self.text is not None:
            provided_fields += 1
        if self.image:
            provided_fields += 1
        if self.file:
            provided_fields += 1
        if self.audio:
            provided_fields += 1
        if self.video:
            provided_fields += 1

        if provided_fields != 1:
            raise ValueError("Exactly one of 'text', 'image', 'file', 'audio', or 'video' must be provided.")
        return self

    def __repr__(self):
        return f"ContentItem({self.model_dump()})"

    def get_type_and_value(self) -> Tuple[Literal["text", "image", "file", "audio", "video"], str]:
        ((t, v),) = self.model_dump().items()
        assert t in ("text", "image", "file", "audio", "video")
        return t, v

    @property
    def type(self) -> Literal["text", "image", "file", "audio", "video"]:
        t, v = self.get_type_and_value()
        return t

    @property
    def value(self) -> str:
        t, v = self.get_type_and_value()
        return v


class Message(BaseModel):
    """Base message class."""

    role: str
    """The role of the messages author (system, user, assistant or tool)."""

    content: Optional[Union[str, List[Union[str, dict]], List[ContentItem]]] = None
    """Content of the message."""

    name: Optional[str] = None
    """An optional name."""

    tool_call_id: Optional[str] = None
    """Tool call that this message is responding to."""

    tool_calls: Optional[List[Dict[str, Any]]] = None
    """The tool calls generated by the model, such as function calls."""

    reasoning_content: Optional[str] = None
    """Thinking content from the model."""

    extra: Optional[dict] = None

    model_config = ConfigDict(extra="allow", populate_by_name=True, arbitrary_types_allowed=True)

    def __repr__(self):
        return f"Message({self.model_dump()})"

    @field_validator("role")
    @classmethod
    def role_checker(cls, value: str) -> str:
        if value not in [USER, ASSISTANT, SYSTEM, TOOL]:
            raise ValueError(f"{value} must be one of {','.join([USER, ASSISTANT, SYSTEM, TOOL])}")
        return value

    def to_dict(self, **kwargs: Any) -> Dict[str, Any]:
        message_dict = self.model_dump(**kwargs)
        message_dict = {
            k: v for k, v in message_dict.items() if v is not None and not (isinstance(v, (list, dict)) and len(v) == 0)
        }
        return message_dict

    def to_openai(self) -> Dict[str, Any]:
        """Format a message into the format expected by OpenAI."""
        message_dict: Dict[str, Any] = {
            "role": self.role,
            "content": self.content,
            "name": self.name,
            "tool_call_id": self.tool_call_id,
            "tool_calls": self.tool_calls,
        }
        message_dict = {k: v for k, v in message_dict.items() if v is not None}

        # OpenAI expects the tool_calls to be None if empty, not an empty list
        if self.tool_calls is not None and len(self.tool_calls) == 0:
            message_dict["tool_calls"] = None

        return message_dict

    # def to_function_call_dict(self) -> Dict[str, Any]:
    #     return {
    #         "content": self.content,
    #         "tool_call_id": self.tool_call_id,
    #         "tool_name": self.tool_name,
    #         "tool_args": self.tool_args,
    #         "tool_call_error": self.tool_call_error,
    #         "metrics": self.metrics,
    #         "created_at": self.created_at,
    #     }