"""Agent Types and Schemas

Defines data models and type aliases used by agents.
"""

from __future__ import annotations

from collections.abc import Awaitable, Callable
from dataclasses import dataclass
from typing import Annotated, Any, Literal, TypeVar

from pydantic import BaseModel, Field

T = TypeVar("T", bound=BaseModel)

# Reasoning effort levels for thinking models (Gemini, Claude, etc.)
# None = no reasoning (default for most agents)
# "low"/"medium"/"high" = increasing reasoning depth
ReasoningEffort = Literal["low", "medium", "high"] | None


@dataclass
class AgentConfig:
    """Configuration for an agent's LLM backend.

    All providers/models supported by LiteLLM are accepted.
    Model name is built as {provider}/{model} for LiteLLM.
    """

    provider: str  # Provider name (e.g., "anthropic", "openai", "gemini")
    model: str  # Model name (e.g., "claude-haiku-4-5-20251001")
    api_key: str  # API key for the provider
    api_base: str | None = None  # Optional custom URL (Ollama, self-hosted)
    tier: str | None = None  # Agent tier (e.g., "review", "utility") for tier-specific config


class TokenUsage(BaseModel):
    """Token usage statistics from LLM call.

    Supports arithmetic operations for easy aggregation:
        total = usage1 + usage2
        total += usage3
    """

    input_tokens: Annotated[int, Field(description="Number of input tokens")]
    output_tokens: Annotated[int, Field(description="Number of output tokens")]
    total_tokens: Annotated[int, Field(description="Total tokens used")]
    cached_tokens: Annotated[
        int, Field(description="Number of cached tokens (prompt cache hits)", default=0)
    ]
    model: Annotated[str, Field(description="Model that was used")]
    exploration_cost: Annotated[
        TokenUsage | None,
        Field(description="Token usage from exploration (if any)", default=None),
    ] = None

    def __add__(self, other: TokenUsage) -> TokenUsage:
        """Add two TokenUsage objects together.

        Args:
            other: Another TokenUsage to add

        Returns:
            New TokenUsage with summed values. Uses self.model for the result.
        """
        # Aggregate exploration costs if present
        exploration = None
        if self.exploration_cost and other.exploration_cost:
            exploration = self.exploration_cost + other.exploration_cost
        elif self.exploration_cost:
            exploration = self.exploration_cost
        elif other.exploration_cost:
            exploration = other.exploration_cost

        return TokenUsage(
            input_tokens=self.input_tokens + other.input_tokens,
            output_tokens=self.output_tokens + other.output_tokens,
            total_tokens=self.total_tokens + other.total_tokens,
            cached_tokens=self.cached_tokens + other.cached_tokens,
            model=self.model,
            exploration_cost=exploration,
        )

    def __iadd__(self, other: TokenUsage) -> TokenUsage:
        """In-place addition (+=) for TokenUsage.

        Note: Returns a new object since Pydantic models are immutable by default.
        """
        return self + other

    @classmethod
    def zero(cls, model: str = "") -> TokenUsage:
        """Create a zero-valued TokenUsage for accumulation.

        Args:
            model: Model name to use

        Returns:
            TokenUsage with all zero values
        """
        return cls(
            input_tokens=0,
            output_tokens=0,
            total_tokens=0,
            cached_tokens=0,
            model=model,
        )

    @classmethod
    def from_litellm_usage(cls, usage: Any, model: str) -> TokenUsage:
        """Create TokenUsage from LiteLLM response usage object.

        LiteLLM uses different field names:
        - prompt_tokens -> input_tokens
        - completion_tokens -> output_tokens

        Args:
            usage: LiteLLM Usage object (or None)
            model: Model name to use

        Returns:
            TokenUsage with values from the LiteLLM response
        """
        if not usage:
            return cls.zero(model=model)

        return cls(
            input_tokens=getattr(usage, "prompt_tokens", 0) or 0,
            output_tokens=getattr(usage, "completion_tokens", 0) or 0,
            total_tokens=getattr(usage, "total_tokens", 0) or 0,
            cached_tokens=cls._extract_cached_tokens(usage),
            model=model,
        )

    @staticmethod
    def _extract_cached_tokens(usage: Any) -> int:
        """Extract cached token count from usage object.

        Different providers use different field names for cached tokens.
        """
        if not usage:
            return 0

        # Gemini
        if hasattr(usage, "cached_content_token_count") and usage.cached_content_token_count:
            return usage.cached_content_token_count

        # Anthropic/OpenRouter
        if hasattr(usage, "cache_read_input_tokens") and usage.cache_read_input_tokens:
            return usage.cache_read_input_tokens

        # OpenAI
        if hasattr(usage, "prompt_tokens_details") and usage.prompt_tokens_details:
            return getattr(usage.prompt_tokens_details, "cached_tokens", 0) or 0

        return 0


class AgentResponse[T: BaseModel](BaseModel):
    """Response from an agent including output and metadata."""

    output: Annotated[T, Field(description="The parsed output from the agent")]
    metadata: Annotated[TokenUsage, Field(description="Token usage information")]


class ToolResult(BaseModel):
    """Result from processing a tool call.

    Used by tool handlers to return content and optional token usage.
    """

    content: Annotated[str, Field(description="The tool result content to send back to the LLM")]
    usage: Annotated[
        TokenUsage | None,
        Field(
            description="Token usage from the tool execution (e.g., explorer calls)", default=None
        ),
    ] = None


# Type alias for async tool handlers
# Takes a tool_call object from LiteLLM response, returns ToolResult
ToolHandler = Callable[[Any], Awaitable[ToolResult]]
