"""
Agent models for CMDOP SDK.

Provides models for AI agent execution and streaming events.
"""

from __future__ import annotations

from enum import Enum
from typing import Any, Generic, TypeVar

from pydantic import BaseModel, ConfigDict, Field

T = TypeVar("T")


class AgentType(str, Enum):
    """Type of AI agent to execute."""

    CHAT = "chat"
    """Simple chat agent without terminal access."""

    TERMINAL = "terminal"
    """Terminal agent with PTY access for command execution."""

    COMMAND = "command"
    """Command execution agent."""

    ROUTER = "router"
    """Router agent that delegates to specialist agents."""

    PLANNER = "planner"
    """Planner agent for complex multi-step tasks."""


class AgentEventType(str, Enum):
    """Type of streaming event from agent execution."""

    TOKEN = "token"
    """Token generated by LLM."""

    TOOL_START = "tool_start"
    """Tool execution started."""

    TOOL_END = "tool_end"
    """Tool execution completed."""

    THINKING = "thinking"
    """Agent is processing/reasoning."""

    ERROR = "error"
    """Error occurred during execution."""

    HANDOFF = "handoff"
    """Agent handed off to another agent."""

    CANCELLED = "cancelled"
    """Agent run was cancelled."""



class AgentUsage(BaseModel):
    """Token usage statistics for agent execution."""

    model_config = ConfigDict(extra="forbid")

    prompt_tokens: int = Field(default=0, ge=0)
    """Number of tokens in the prompt."""

    completion_tokens: int = Field(default=0, ge=0)
    """Number of tokens in the completion."""

    total_tokens: int = Field(default=0, ge=0)
    """Total tokens used."""


class AgentToolResult(BaseModel):
    """Result of a tool execution during agent run."""

    model_config = ConfigDict(extra="forbid")

    tool_name: str
    """Name of the tool that was executed."""

    tool_call_id: str = ""
    """Unique ID for this tool call."""

    success: bool
    """Whether the tool execution succeeded."""

    result: str = ""
    """Result output from the tool."""

    error: str = ""
    """Error message if tool failed."""

    duration_ms: int = Field(default=0, ge=0)
    """Execution duration in milliseconds."""


class AgentRunOptions(BaseModel):
    """Options for agent execution."""

    model_config = ConfigDict(extra="forbid")

    model: str | None = None
    """LLM model to use (e.g., 'openai/gpt-4o')."""

    max_turns: int = Field(default=10, ge=1, le=50)
    """Maximum conversation turns."""

    max_retries: int = Field(default=3, ge=0, le=10)
    """Maximum retries on failure."""

    timeout_seconds: int = Field(default=300, ge=10, le=600)
    """Timeout in seconds."""

    stream_events: bool = False
    """Whether to stream events in real-time."""

    def to_options_map(self) -> dict[str, str]:
        """Convert to proto options map."""
        opts: dict[str, str] = {}
        if self.model:
            opts["model"] = self.model
        opts["max_turns"] = str(self.max_turns)
        opts["max_retries"] = str(self.max_retries)
        opts["timeout_seconds"] = str(self.timeout_seconds)
        return opts


class AgentStreamEvent(BaseModel):
    """Real-time event from agent execution."""

    model_config = ConfigDict(extra="forbid")

    request_id: str
    """Request ID for correlation."""

    type: AgentEventType
    """Type of event."""

    payload: Any = None
    """Event payload data (parsed from JSON)."""

    timestamp: int = 0
    """Unix timestamp in milliseconds."""

    @property
    def payload_data(self) -> dict[str, Any]:
        """Get payload as dict (for backwards compatibility)."""
        import json

        if isinstance(self.payload, dict):
            return self.payload
        if not self.payload:
            return {}
        try:
            return json.loads(self.payload)
        except json.JSONDecodeError:
            return {"raw": self.payload}


class AgentResult(BaseModel, Generic[T]):
    """Result of agent execution.

    When using structured output (output_model parameter), the result
    will have the `data` field populated with the parsed Pydantic model.

    Example:
        >>> class Answer(BaseModel):
        ...     value: int
        ...     explanation: str
        >>>
        >>> result = await client.agent.run("2+2?", output_model=Answer)
        >>> print(result.data.value)  # 4
    """

    model_config = ConfigDict(extra="forbid")

    request_id: str
    """Request ID for correlation."""

    success: bool
    """Whether the agent execution succeeded."""

    text: str = ""
    """Agent's text response."""

    error: str = ""
    """Error message if failed."""

    tool_results: list[AgentToolResult] = Field(default_factory=list)
    """Results from tool executions."""

    usage: AgentUsage = Field(default_factory=AgentUsage)
    """Token usage statistics."""

    duration_ms: int = Field(default=0, ge=0)
    """Total execution duration in milliseconds."""

    data: T | None = None
    """Structured output data (if output_model was provided)."""

    output_json: str = ""
    """Raw JSON string of structured output."""

    @property
    def duration_seconds(self) -> float:
        """Duration in seconds."""
        return self.duration_ms / 1000.0


class AgentRunRequest(BaseModel):
    """Request to run an AI agent."""

    model_config = ConfigDict(extra="forbid")

    prompt: str
    """The prompt/question for the agent."""

    agent_type: AgentType = AgentType.CHAT
    """Type of agent to run."""

    options: AgentRunOptions = Field(default_factory=AgentRunOptions)
    """Execution options."""
