"""Type definitions for web search endpoint."""

from typing import Optional

from openai import BaseModel as OpenAIBaseModel
from pydantic import BaseModel, Field


class UserLocation(BaseModel):
    """User location for web search localization."""

    type: str = "approximate"
    """Location type (required by Anthropic API)."""

    country: Optional[str] = None
    """ISO 3166-1 alpha-2 country code (e.g., "US", "GB", "KR")."""

    city: Optional[str] = None
    """City name."""

    region: Optional[str] = None
    """Region/state name."""

    timezone: Optional[str] = None
    """Timezone (e.g., "Asia/Seoul", "America/New_York")."""


class SearchCitation(OpenAIBaseModel):
    """Source citation from web search."""

    title: str
    """Title of the source page."""

    url: str
    """URL of the source."""

    cited_text: Optional[str] = None
    """Text excerpt that was cited."""


class SearchUsage(OpenAIBaseModel):
    """Token usage for web search."""

    input_tokens: int
    """Number of input tokens used."""

    output_tokens: int
    """Number of output tokens generated."""


class WebSearchRequest(BaseModel):
    """Request for web search."""

    query: str = Field(..., min_length=1, max_length=2000)
    """Search query."""

    model: str = "claude-3-5-haiku-20241022"
    """Model to use for search. Defaults to claude-3-5-haiku."""

    max_tokens: int = Field(default=1024, ge=1, le=4096)
    """Maximum tokens in response."""

    max_searches: int = Field(default=5, ge=1, le=10)
    """Maximum number of web searches to perform."""

    allowed_domains: Optional[list[str]] = None
    """Restrict search to these domains only."""

    blocked_domains: Optional[list[str]] = None
    """Exclude these domains from search."""

    user_location: Optional[UserLocation] = None
    """User location for localized results."""


class WebFetchRequest(BaseModel):
    """Request to fetch and analyze URL content."""

    url: str = Field(..., min_length=1)
    """URL to fetch and analyze."""

    prompt: str = "Summarize this page"
    """Prompt for analyzing the content."""

    model: str = "claude-3-5-haiku-20241022"
    """Model to use. Defaults to claude-3-5-haiku."""

    max_tokens: int = Field(default=1024, ge=1, le=4096)
    """Maximum tokens in response."""


class WebSearchResponse(OpenAIBaseModel):
    """Response from web search."""

    id: str
    """Unique response ID."""

    content: str
    """AI-generated response with search results."""

    citations: list[SearchCitation] = []
    """List of source citations."""

    model: str
    """Model used for the search."""

    usage: SearchUsage
    """Token usage information."""

    stop_reason: Optional[str] = None
    """Reason the response stopped."""
