"""
Models for the Bedrock agent.
"""

from enum import StrEnum
from typing import Any

from pydantic import BaseModel, Field

#######################
# Agent Event Classes #
#######################


class AgentEventTypes(StrEnum):
    """Enumeration of possible agent event types."""

    START = "start"
    STOP = "stop"
    THINKING = "thinking"
    ANSWER = "answer"
    TOOL_USE_START = "tool_use_start"
    TOOL_USE_STOP = "tool_use_stop"
    DELEGATION_START = "delegation_start"
    DELEGATION_STOP = "delegation_stop"


class AgentEvent(BaseModel):
    """Base class for all agent events.

    :param type: The type of event
    :param content: The content of the event
    """

    agent_name: str
    type: AgentEventTypes
    content: Any

    model_config = {"use_enum_values": True}

    def __str__(self):
        return f"{self.agent_name} [{self.type.capitalize()}] {self.content}"

    def __repr__(self):
        return self.__str__()


class AgentStartEvent(AgentEvent):
    """Event representing an agent's start."""

    type: AgentEventTypes = AgentEventTypes.START
    content: str | None = None

    def __str__(self):
        return f"{self.agent_name} [{self.type.capitalize()}]"

    def __repr__(self):
        return self.__str__()


class AgentStopEvent(AgentEvent):
    """Event representing an agent's stop."""

    type: AgentEventTypes = AgentEventTypes.STOP
    content: str | None = None

    def __str__(self):
        return f"{self.agent_name} [{self.type.capitalize()}]"

    def __repr__(self):
        return self.__str__()


#######################
# Tool Related Models #
#######################


class ToolUseData(BaseModel):
    """Data model for tool usage information.

    :param name: Name of the tool being used
    :param parameters: Parameters passed to the tool
    :param result: Result returned by the tool
    """

    tool_name: str
    parameters: dict[str, Any] = Field(default_factory=dict)
    result: str | dict | None = None

    def __str__(self):
        return f"{self.tool_name} ({self.parameters}) -> {self.result}"

    def __repr__(self):
        return self.__str__()


class ToolUseStartEvent(AgentEvent):
    """Event representing an agent's use of a tool."""

    type: AgentEventTypes = AgentEventTypes.TOOL_USE_START
    content: ToolUseData

    def __str__(self):
        return f"{self.agent_name} [{self.type.capitalize()}] {self.content}"

    def __repr__(self):
        return self.__str__()


class ToolUseStopEvent(AgentEvent):
    """Event representing an agent's tool use stop."""

    type: AgentEventTypes = AgentEventTypes.TOOL_USE_STOP
    content: ToolUseData


#######################
# Text Related Models #
#######################
class ThinkingEvent(AgentEvent):
    """Event representing an agent's thinking process."""

    type: AgentEventTypes = AgentEventTypes.THINKING
    content: str

    def __str__(self):
        return f"{self.agent_name} [{self.type.capitalize()}] {self.content}"

    def __repr__(self):
        return self.__str__()


class AnswerEvent(AgentEvent):
    """Event representing an agent's answer."""

    type: AgentEventTypes = AgentEventTypes.ANSWER
    content: str

    def __str__(self):
        return f"{self.agent_name} [{self.type.capitalize()}] {self.content}"

    def __repr__(self):
        return self.__str__()


#############################
# Delegation Related Models #
#############################


class DelegateData(BaseModel):
    """Data model for delegation information.

    :param agent_name: Name of the delegate agent
    :param query: Query passed to the delegate
    :param answer: Answer returned by the delegate
    :param events: List of events generated by the delegate agent
    """

    agent_name: str
    query: str
    answer: str | None = None
    events: list[AgentEvent] = Field(
        default_factory=list,
        description="List of events generated by the delegate agent.",
    )

    def __str__(self):
        events_str = (
            "\n" + "\n".join(f"\t>{str(event)}" for event in self.events)
            if self.events
            else ""
        )
        return f"{self.agent_name} ({self.query}) -> {self.answer}{events_str}"

    def __repr__(self):
        return self.__str__()


class DelegateStartEvent(AgentEvent):
    """Event representing an agent delegating a task to another agent."""

    type: AgentEventTypes = AgentEventTypes.DELEGATION_START
    content: DelegateData


class DelegateStopEvent(AgentEvent):
    """Event representing an agent's delegation stop."""

    type: AgentEventTypes = AgentEventTypes.DELEGATION_STOP
    content: DelegateData


###################
# Response Models #
###################


class AgentResponse(BaseModel):
    """Model representing a complete response from an agent.

    :param answer: The final answer from the agent
    :param events: List of events that occurred during the agent's execution
    """

    answer: str
    events: list[AgentEvent]

    model_config = {"use_enum_values": True}
