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

from __future__ import annotations
from dataclasses import asdict, dataclass, field, replace
from typing import Any, Sequence
import json


@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]

    @classmethod
    def load(cls, data: dict[str, Any]) -> Message:
        return cls(
            role=data["role"],
            content=data["content"]
            if isinstance(data["content"], str)
            else [MessageContent.load(content) for content in data["content"]],
        )

    def dump(self):
        return {
            "role": self.role,
            "content": self.content
            if isinstance(self.content, str)
            else [content.dump() for content in self.content],
        }


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

    role: str = field(init=False, default="system")


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

    role: str = field(init=False, default="user")


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

    # The resource usage for the message.
    role: str = field(init=False, default="assistant")

    # The resource usage for the message.
    usage: Usage = field(default_factory=lambda: Usage())


@dataclass(frozen=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(frozen=True)
class TextBlock(MessageContent):
    """A block of text content."""

    text: str

    def append(self, text: str) -> TextBlock:
        """Append text to the block.

        Args:
            text: The text to append.

        Returns:
            A new TextBlock instance with the appended text.
        """
        return replace(self, text=self.text + text)


@dataclass(frozen=True)
class DataBlock(MessageContent):
    """A block of structured data."""

    data: dict[str, Any]


@dataclass(kw_only=True, frozen=True)
class ToolUseChunk(MessageContent):
    """A chunk reported by the tool during execution."""

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

    # The name of the tool this is from.
    tool_name: str

    # Associated data.
    content: Any


@dataclass(kw_only=True, frozen=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

    def append_input(self, chunk: str) -> ToolCall:
        """Append a chunk to the input data.

        Args:
            chunk: The chunk of text to append from the stream.

        Returns:
            A new ToolCall instance with the appended input.
        """
        assert isinstance(self.input, str)
        return replace(self, input=self.input + chunk)

    def load_input(self) -> ToolCall:
        """Load the input string as JSON.

        Returns:
            A new ToolCall instance with the input loaded as JSON.
        """
        assert isinstance(self.input, str)
        try:
            data = json.loads(self.input)
        except json.JSONDecodeError:
            data = {}
        return replace(self, input=data)


@dataclass(kw_only=True, frozen=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


@dataclass(kw_only=True)
class Usage:
    """Track resource usage for a message."""

    # The number of input tokens used to generate the message.
    input_tokens: int = 0

    # The number of output tokens used to generate the message.
    output_tokens: int = 0

    # The number of tokens read from the cache.
    cache_read_tokens: int = 0

    # The number of tokens written to the cache.
    cache_write_tokens: int = 0

    # The time taken to generate the message.
    generation_time: float = 0.0
