"""Type definitions for OCR endpoint."""

from typing import Literal, Optional

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

# Model quality for auto-selection
ModelQuality = Literal["fast", "balanced", "best"]


class OCRRequest(BaseModel):
    """Request for OCR text extraction."""

    image: Optional[str] = None
    """Base64-encoded image data."""

    image_url: Optional[str] = None
    """URL of the image to extract text from."""

    model: Optional[str] = None
    """Model to use for OCR. Defaults to server default."""

    model_quality: Optional[ModelQuality] = None
    """Model quality preset for auto-selection: 'fast' (cheapest), 'balanced', 'best'.
    Ignored if explicit model is specified."""

    mode: Literal["tiny", "small", "base", "gundam"] = "base"
    """OCR mode affecting quality and speed."""

    fetch_image: bool = True
    """Whether to fetch image from URL on server side."""

    @model_validator(mode="after")
    def validate_image_source(self) -> "OCRRequest":
        """Ensure at least one image source is provided."""
        if not self.image and not self.image_url:
            raise ValueError("Either 'image' or 'image_url' must be provided")
        return self


class OCRResponse(OpenAIBaseModel):
    """Response from OCR extraction."""

    text: str
    """Extracted text content."""

    model: str
    """Model used for extraction."""

    cost_usd: float
    """Cost of the request in USD."""

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

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