from typing import Any

from pydantic import (
    GetCoreSchemaHandler,
    GetJsonSchemaHandler,
    SerializationInfo,
    TypeAdapter,
    ValidationInfo,
    ValidatorFunctionWrapHandler,
)
from pydantic_core import CoreSchema, core_schema

from .content import TextContent, ToolResultContent, ToolUseContent, content_factory


class Message:
    """A class representing a message in a conversation with an LLM.

    This class handles messages for both OpenAI and Anthropic formats, providing
    conversion methods between different message formats and validation logic.

    Attributes:
        role: The role of the message sender
        content: The content of the message, which can include text and tool interactions
        stop_reason: The reason the LLM stopped generating (only for assistant messages).
            Provider-specific values:
            - Anthropic: 'end_turn', 'max_tokens', 'stop_sequence', 'tool_use', 'pause_turn', 'refusal'
            - OpenAI Chat Completions: 'stop', 'length', 'tool_calls', 'content_filter', 'function_call'
            - OpenAI Responses: 'completed', 'max_output_tokens', 'content_filter', etc.
    """

    __slots__ = ("role", "content", "stop_reason")

    def __init__(self, role: Any, content: Any, stop_reason: str | None = None) -> None:
        """Initialize a Message instance.

        Args:
            role: The role of the message sender
            content: The content of the message
            stop_reason: The reason the LLM stopped generating (optional, only for assistant messages)
        """
        object.__setattr__(self, "role", role)
        object.__setattr__(self, "content", content)
        object.__setattr__(self, "stop_reason", stop_reason)

    def __str__(self) -> str:
        if self.stop_reason:
            return f"Message(role={self.role}, content={self.content}, stop_reason={self.stop_reason})"
        return f"Message(role={self.role}, content={self.content})"

    def __repr__(self) -> str:
        if self.stop_reason:
            return f"Message(role={self.role}, content={self.content}, stop_reason={self.stop_reason})"
        return f"Message(role={self.role}, content={self.content})"

    def to_openai_responses_input(self) -> list[dict[str, Any]]:
        """Convert the message to OpenAI's responses api expected input format."""
        inputs = []
        message_content = []
        for content_item in self.content:
            if isinstance(content_item, ToolUseContent):
                inputs.append(content_item.to_openai_responses_input())
            elif isinstance(content_item, ToolResultContent):
                inputs.append(content_item.to_openai_responses_input())
            else:
                message_content.append(content_item.to_openai_responses_input(role=self.role))
        if message_content:
            # Role here should only be 'user' or 'assistant'
            inputs.append({"role": self.role, "content": message_content})
        return inputs

    def to_openai_chat_completions_input(self) -> dict[str, Any]:
        """Convert the message to OpenAI's chat completions api expected input format."""
        role = self.role
        # OpenAI chat completions api expects tool calls to be in a separate field in the message
        content = []
        tool_calls = []
        for content_item in self.content:
            if isinstance(content_item, ToolUseContent):
                tool_calls.append(content_item.to_openai_chat_completions_input())
            elif isinstance(content_item, ToolResultContent):
                return content_item.to_openai_chat_completions_input()
            else:
                openai_input = content_item.to_openai_chat_completions_input()
                # Enabling splitting files into multiple pages or chunks.
                if isinstance(openai_input, list):
                    content.extend(openai_input)
                else:
                    content.append(openai_input)
        openai_input = {
            "role": role,
        }
        if len(content):
            openai_input["content"] = content
        if len(tool_calls):
            openai_input["tool_calls"] = tool_calls
        return openai_input

    def to_anthropic_input(self) -> dict[str, Any]:
        """Convert the message to Anthropic's expected input format."""
        content = []
        for content_item in self.content:
            anthropic_input = content_item.to_anthropic_input()
            # Enabling splitting files into multiple pages or chunks.
            if isinstance(anthropic_input, list):
                content.extend(anthropic_input)
            else:
                content.append(anthropic_input)
        # Anthropic doesn't accept the tool role. We must send this under the user role.
        role = self.role
        if role == "tool":
            role = "user"
        return {
            "role": role,
            "content": content,
        }

    def collect_text(self) -> str:
        """Collect all text from the message content."""
        message_text = ""
        for content in self.content:
            if isinstance(content, TextContent):
                message_text += content.text + "\n\n"
        return message_text.strip()

    @classmethod
    def validate(cls, value: ValidatorFunctionWrapHandler, _info: dict | ValidationInfo | None = None) -> "Message":
        """Validate and convert inputs into a Message instance."""
        # Don't recurse if we're already dealing with a Message instance
        if isinstance(value, Message):
            return value
        if isinstance(value, dict):
            role = value.get("role", "user")
            content = value.get("content", None)
            stop_reason = value.get("stop_reason", None)
            if not isinstance(content, list):
                content = [content]
            content = [content_factory(item) for item in content]
            return cls(role=role, content=content, stop_reason=stop_reason)
        return cls.validate(
            {
                "role": "user",
                "content": value,
            }
        )

    @classmethod
    def serialize(cls, value: Any, _info: dict | SerializationInfo | None = None) -> str:
        """Serialize a Message instance into a dictionary format."""
        # When creating a model with fields with File type that are nullable,
        # pydantic will pass None as the value to File.serialize.
        if value is None:
            return None
        if not isinstance(value, cls):
            raise ValueError("Cannot serialize a non-message object.")
        result = {
            "role": value.role,
            "content": value.content,
        }
        # Only include stop_reason if it's set (to avoid breaking existing serialization)
        if value.stop_reason is not None:
            result["stop_reason"] = value.stop_reason
        return result

    @classmethod
    def __get_pydantic_json_schema__(cls, _core_schema: CoreSchema, _handler: GetJsonSchemaHandler) -> dict[str, Any]:
        """Defines what this type should be in openapi.json."""
        # https://docs.pydantic.dev/2.8/errors/usage_errors/#custom-json-schema
        json_schema = {
            "title": "TimbalMessage",  # This becomes the type name in most generators
            "type": "object",
            "properties": {
                "role": {
                    "type": "string",
                    "enum": ["user", "assistant", "tool", "system"],
                },
                "content": {
                    "type": "array",
                    "items": {},  # Keep it open/generic for now.
                },
            },
        }
        return json_schema

    @classmethod
    def __get_pydantic_core_schema__(cls, _source: type[Any], _handler: GetCoreSchemaHandler) -> core_schema.CoreSchema:
        """Defines how to serialize this type in the core schema."""
        return core_schema.with_info_plain_validator_function(
            cls.validate,
            serialization=core_schema.plain_serializer_function_ser_schema(
                cls.serialize,
                info_arg=True,
                when_used="always",
            ),
        )


message_model_schema = TypeAdapter(Message).json_schema()
