"""Message types to send between the agent and the LLM."""

from __future__ import annotations
from typing import Any
import json

from anthropic.types import ToolUseBlock

import reflex as rx


class Message(rx.Base):
    """The base message class."""

    # The role of the message (user, system, AI, tool).
    role: str

    # The content of the message.
    content: str


class SystemMessage(Message):
    """A system message."""

    role = "system"


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

    role = "user"


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

    role = "assistant"
    # Token usage for the input that generated this output
    input_tokens: int = 0
    # Token usage to generate this output
    output_tokens: int = 0


class ToolUseMessage(AIMessage):
    """A message to use a tool."""
    # Tool Use ID
    id: str
    # Name of the tool being invoked
    tool_name: str
    # Params/Arguments for the Tool to be called with
    input: Any

    @classmethod
    def from_tool_use_block(cls, tool_use_block: ToolUseBlock, **kwargs) -> ToolUseMessage:
        return cls(
            content=tool_use_block.json(),
            id=tool_use_block.id,
            tool_name=tool_use_block.name,
            input=tool_use_block.input,
            **kwargs,
        )


class ToolResultMessage(UserMessage):
    """A message containing the result of a tool use."""
    # Tool Usage ID
    tool_use_id: str
    # Name of Tool Used
    tool_name: str
    # Inputs for the Tools Used
    tool_args: Any
    # Time taken for the Tool to run
    execution_time: float
    # Was an error generated
    is_error: bool

    @classmethod
    def from_result(cls, tool_use_message: ToolUseMessage, result: Any, execution_time: float) -> ToolResultMessage:
        """Create a tool result message from a tool result.

        Args:
            tool_use_id: The ID of the tool use.
            result: The result of the tool (or an exception).

        Returns:
            A tool result message.
        """
        # The base message data.
        data = {
            "type": "tool_result",
            "content": result,
        }
        is_error = False

        # Check for an exception.
        if isinstance(result, Exception):
            data.update(
                {
                    "content": str(result),
                }
            )
            is_error = True

        # Return the message.
        return cls(
            tool_use_id = tool_use_message.id,
            tool_name = tool_use_message.tool_name,
            tool_args = tool_use_message.input,
            is_error = is_error,
            execution_time = execution_time,
            content=json.dumps([data]))