"""Message types for agent-LLM communication in the FlexAI framework."""

from __future__ import annotations
from dataclasses import asdict, dataclass
from typing import Any, Sequence


@dataclass
class Message:
    """Base class for all message types in the conversation flow."""

    # The role of the message (user or assistant).
    role: str

    # The content of the message.
    content: str | Sequence[MessageContent]


@dataclass
class SystemMessage(Message):
    """A top level system message."""

    def __init__(self, content: str | Sequence[MessageContent]):
        super().__init__("system", content)


@dataclass
class UserMessage(Message):
    """A message sent by a user."""

    def __init__(self, content: str | Sequence[MessageContent]):
        super().__init__("user", content)


@dataclass
class AIMessage(Message):
    """A message generated by the AI."""

    def __init__(self, content: str | Sequence[MessageContent]):
        super().__init__("assistant", content)


@dataclass(kw_only=True)
class MessageContent:
    def dump(self) -> dict[str, Any]:
        """Return the dataclass as a dictionary.

        Returns:
            The dataclass as a dictionary.
        """
        return {"type": self.__class__.__name__, **asdict(self)}

    @classmethod
    def load(cls, data: dict[str, Any]) -> MessageContent:
        """Load the dataclass from a dictionary.

        Args:
            data: The dictionary to load.

        Returns:
            The dataclass instance.

        Raises:
            ValueError: If the message content type is unknown.
        """
        # Check all subclasses of MessageContent.
        for subclass in cls.__subclasses__():
            if subclass.__name__ == data["type"]:
                data = {k: v for k, v in data.items() if k != "type"}
                return subclass(**data)

        raise ValueError(f"Unknown message content type: {data['type']}")


@dataclass(kw_only=True)
class TextBlock(MessageContent):
    """A block of text content."""

    text: str


@dataclass(kw_only=True)
class ToolCall(MessageContent):
    """A tool call message sent by the agent."""

    # A unique identifier for the tool call.
    id: str

    # The name of the tool to call.
    name: str

    # The input parameters for the tool.
    input: Any


@dataclass(kw_only=True)
class ToolResult(MessageContent):
    """A tool result message created after invoking a tool."""

    # The associated tool call identifier.
    tool_call_id: str

    # The result of the tool invocation.
    result: Any

    # The execution time of the tool invocation.
    execution_time: float = 0.0

    # Whether an error occurred during invocation.
    is_error: bool = False
