"""Agent event system for streaming execution updates."""

from typing import Literal, Any
from pydantic import BaseModel

from yier_agents.src.message import Message


class AgentEvent(BaseModel):
    """Base class for all agent events."""
    type: str
    session_id: str
    timestamp: float  # Unix timestamp


class AgentStartEvent(AgentEvent):
    """Emitted when agent starts execution."""
    type: Literal["agent_start"] = "agent_start"
    user_message: str
    iteration: int = 0


class AgentEndEvent(AgentEvent):
    """Emitted when agent completes execution."""
    type: Literal["agent_end"] = "agent_end"
    finish_reason: Literal["stop", "tool_calls", "length", "max_iterations"]
    total_iterations: int


class LLMStartEvent(AgentEvent):
    """Emitted when LLM call starts."""
    type: Literal["llm_start"] = "llm_start"
    iteration: int


class LLMEndEvent(AgentEvent):
    """Emitted when LLM call completes."""
    type: Literal["llm_end"] = "llm_end"
    iteration: int
    message: Message
    finish_reason: Literal["stop", "tool_calls", "length"]
    usage: int | None = None


class ToolCallStartEvent(AgentEvent):
    """Emitted when tool execution starts."""
    type: Literal["tool_call_start"] = "tool_call_start"
    iteration: int
    tool_name: str
    tool_call_id: str
    arguments: dict[str, Any]


class ToolCallEndEvent(AgentEvent):
    """Emitted when tool execution completes."""
    type: Literal["tool_call_end"] = "tool_call_end"
    iteration: int
    tool_name: str
    tool_call_id: str
    result: str
    is_error: bool = False


class MessageCompactEvent(AgentEvent):
    """Emitted when messages are being compacted due to token limit."""
    type: Literal["message_compact"] = "message_compact"
    iteration: int
    original_count: int
    compacted_count: int


class ErrorEvent(AgentEvent):
    """Emitted when an error occurs."""
    type: Literal["error"] = "error"
    iteration: int
    error_message: str
    error_type: str


class MemoryEvent(AgentEvent):
    """Emitted when session memory is loaded or saved."""
    type: Literal["memory"] = "memory"
    action: Literal["load", "save", "clear"]
    message_count: int


# Union type for all events
AgentEventType = (
    AgentStartEvent
    | AgentEndEvent
    | LLMStartEvent
    | LLMEndEvent
    | ToolCallStartEvent
    | ToolCallEndEvent
    | MessageCompactEvent
    | ErrorEvent
    | MemoryEvent
)
