import json
import uuid
from abc import ABC
from typing import Any, Literal, Self
from pydantic import BaseModel, ConfigDict, Field, field_validator
from .attachment import Attachment

class ChatMessage(BaseModel, ABC):
    model_config = ConfigDict(
        arbitrary_types_allowed=True,
        validate_assignment=True,
    )
    id: str = Field(default_factory=lambda: str(uuid.uuid4()))

class SystemMessage(ChatMessage):
    model_config = ConfigDict(json_schema_extra={
        "required": ["content", "role"]
    })

    content: str
    role: Literal["system"] = "system"

class ToolMessage(ChatMessage):
    model_config = ConfigDict(json_schema_extra={
        "required": ["call_id", "name", "arguments", "result", "error", "role", "metadata"]
    })

    call_id: str
    name: str
    arguments: dict[str, Any]
    result: str | None = None
    error: str | None = None
    role: Literal["tool"] = "tool"
    metadata: dict[str, Any] = Field(default_factory=dict)

    @field_validator("result", mode="before")
    def validate_result(cls, v: Any) -> Any:
        if v is None: return v
        if isinstance(v, str): return v
        return json.dumps(v, ensure_ascii=False)

    @property
    def is_complete(self) -> bool:
        return self.result is not None or self.error is not None

    def with_result(self, result: str | None, error: str | None) -> ToolMessage:
        return ToolMessage(
            call_id=self.call_id,
            name=self.name,
            arguments=self.arguments,
            result=result,
            error=error)

class AssistantMessage(ChatMessage):
    class ToolCall(BaseModel):
        id: str
        name: str
        arguments: dict[str, Any]

    class Usage(BaseModel):
        input_tokens: int
        output_tokens: int
        total_tokens: int

        @classmethod
        def default(cls) -> Self:
            return cls(
                input_tokens=0,
                output_tokens=0,
                total_tokens=0)

    model_config = ConfigDict(json_schema_extra={
        "required": ["content", "reasoning_content", "tool_calls", "audio", "images", "usage", "role"]
    })

    content: str | None = None
    reasoning_content: str | None = None
    tool_calls: list[ToolCall] | None = None
    usage: Usage | None = None
    role: Literal["assistant"] = "assistant"

    def get_incomplete_tool_messages(self) -> list[ToolMessage] | None:
        """
        Get a incomplete tool message from the assistant message.
        The returned tool message is incomplete,
        which means it only contains the tool call id, name and arguments.
        Returns None if there is no tool call in the assistant message.
        """
        if self.tool_calls is None: return None
        results: list[ToolMessage] = []
        for tool_call in self.tool_calls:
            results.append(ToolMessage(
                call_id=tool_call.id,
                name=tool_call.name,
                arguments=tool_call.arguments,
                result=None,
                error=None))
        return results

class UserMessage(ChatMessage):
    model_config = ConfigDict(json_schema_extra={
        "required": ["content", "role"]
    })

    content: str
    attachments: list[Attachment] | None = None
    role: Literal["user"] = "user"

__all__ = [
    "ChatMessage",
    "SystemMessage",
    "UserMessage",
    "AssistantMessage",
    "ToolMessage",
]
