"""Primitive types and constants for LLM interactions.

This module provides low-level types for LLM communication that have NO dependency
on the Document class. These types are used by internal modules (document_store,
observability) that need LLM access but cannot depend on Documents.

All types are frozen Pydantic models for immutability and JSON serialization.
"""

from collections.abc import Mapping
from enum import StrEnum
from typing import Any, Literal, cast

from pydantic import Base64Bytes, BaseModel, ConfigDict, field_validator

# Token count per image/PDF (per CLAUDE.md §3.6)
TOKENS_PER_IMAGE = 1080


type ModelName = (
    Literal[
        # Core models
        "gemini-3-pro",
        "gpt-5.1",
        # Small models
        "gemini-3-flash",
        "gpt-5-mini",
        "grok-4.1-fast",
        # Search models
        "gemini-3-flash-search",
        "gpt-5-mini-search",
        "grok-4.1-fast-search",
        "sonar-pro-search",
    ]
    | str
)
"""Type-safe model name with IDE autocompletion for common models.

Literal[...] | str provides autocomplete for known models while accepting any string.
Model availability depends on your LiteLLM proxy configuration and provider access.
"""


class Role(StrEnum):
    """Message role in conversation."""

    SYSTEM = "system"
    USER = "user"
    ASSISTANT = "assistant"


class TextContent(BaseModel):
    """Plain text content part."""

    model_config = ConfigDict(frozen=True)

    type: Literal["text"] = "text"
    text: str


class ImageContent(BaseModel):
    """Image content with binary data.

    Uses Base64Bytes for automatic base64 encoding/decoding in JSON serialization.
    """

    model_config = ConfigDict(frozen=True)

    type: Literal["image"] = "image"
    data: Base64Bytes
    mime_type: Literal["image/jpeg", "image/png", "image/gif", "image/webp"]


class PDFContent(BaseModel):
    """PDF document content with binary data.

    Uses Base64Bytes for automatic base64 encoding/decoding in JSON serialization.
    """

    model_config = ConfigDict(frozen=True)

    type: Literal["pdf"] = "pdf"
    data: Base64Bytes


ContentPart = TextContent | ImageContent | PDFContent
"""Union of all content part types for polymorphic handling."""


class CoreMessage(BaseModel):
    """A single message in a conversation.

    Content can be:
    - str: Plain text (converted to TextContent internally)
    - ContentPart: Single content part (text, image, or PDF)
    - tuple[ContentPart, ...]: Multiple content parts (multimodal message)
    """

    model_config = ConfigDict(frozen=True)

    role: Role
    content: str | ContentPart | tuple[ContentPart, ...]


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

    model_config = ConfigDict(frozen=True)

    prompt_tokens: int  # Input tokens consumed by prompt and context
    completion_tokens: int  # Output tokens generated by the model
    total_tokens: int  # prompt_tokens + completion_tokens
    cached_tokens: int = 0  # Prompt tokens served from provider cache
    reasoning_tokens: int = 0  # Tokens used for internal model reasoning


class ModelOptions(BaseModel):
    r"""Configuration options for LLM generation requests.

    All fields are optional with sensible defaults. Extra fields are forbidden.

    Non-obvious behaviors:
    - cache_ttl: string format ("60s", "5m", "1h") or None to disable. Default "300s".
    - service_tier: only OpenAI models honor this; other providers silently ignore it.
    - stop: accepts a single string or list; coerced to tuple. Max 4 by most providers.
    - usage_tracking: when True (default), injects {"usage": {"include": True}} into extra_body.
    - extra_body: merged with usage_tracking dict if both are set.
    """

    model_config = ConfigDict(frozen=True, extra="forbid")

    temperature: float | None = None  # Leave unset (provider decides) unless you have a specific reason
    system_prompt: str | None = None  # Prepended to messages; substitutor instructions appended when active
    search_context_size: Literal["low", "medium", "high"] | None = None
    reasoning_effort: Literal["low", "medium", "high"] | None = None
    retries: int = 3  # Cache auto-disabled on retry; degeneration detection triggers auto-retry
    retry_delay_seconds: int = 20  # Fixed delay between retries
    timeout: int = 600
    cache_ttl: str | None = "300s"
    service_tier: Literal["auto", "default", "flex", "scale", "priority"] | None = None
    max_completion_tokens: int | None = None  # Defaults to provider behavior (~30K typical)
    stop: tuple[str, ...] | None = None
    verbosity: Literal["low", "medium", "high"] | None = None
    stream: bool = False
    usage_tracking: bool = True
    user: str | None = None
    metadata: Mapping[str, str] | None = None
    extra_body: Mapping[str, Any] | None = None

    @field_validator("stop", mode="before")
    @classmethod
    def _coerce_stop(cls, v: Any) -> tuple[str, ...] | None:
        if v is None:
            return None
        if isinstance(v, str):
            return (v,)
        if isinstance(v, (list, tuple)):
            return tuple(str(s) for s in cast(list[Any], v))
        return v

    def to_openai_completion_kwargs(self) -> dict[str, Any]:
        """Convert options to OpenAI API completion parameters.

        Only includes non-None values. Framework-only fields (system_prompt,
        retries, retry_delay_seconds, cache_ttl, stream) are excluded.
        """
        kwargs: dict[str, Any] = {
            "timeout": self.timeout,
            "extra_body": dict(self.extra_body) if self.extra_body else {},
        }

        # Direct 1:1 field mappings (field name == API kwarg name)
        for attr in ("temperature", "max_completion_tokens", "reasoning_effort", "service_tier", "verbosity", "user"):
            if (v := getattr(self, attr)) is not None:
                kwargs[attr] = v

        # Fields needing transformation
        if self.stop is not None:
            kwargs["stop"] = list(self.stop)
        if self.search_context_size:
            kwargs["web_search_options"] = {"search_context_size": self.search_context_size}
        if self.metadata:
            kwargs["metadata"] = dict(self.metadata)
        if self.usage_tracking:
            kwargs["extra_body"]["usage"] = {"include": True}
            kwargs["stream_options"] = {"include_usage": True}

        return kwargs
