from typing import List, Optional, Dict, Any
from pydantic import BaseModel, Field, field_validator


class VideoRAGAnswer(BaseModel):
    """Default structured answer from VideoRAG."""

    answer: str = Field(..., description="Answer: yes/no/cannot_answer")
    confidence_score: float = Field(
        ...,
        ge=0.0,
        le=1.0,
        description=(
            "Confidence score 0.0-1.0 based on the evidence provided such as "
            "subtitles, music, people, etc. stored into lancedb metadata table"
        ),
    )


class DetailedAnswer(BaseModel):
    """Detailed answer with reasoning and evidence.

    Public API - available for custom output schemas in VideoRAGSession.
    """

    answer: str = Field(..., description="yes/no/cannot_answer")
    confidence_score: float = Field(..., ge=0.0, le=1.0)
    reasoning: str = Field(..., description="Explanation of the answer")
    evidence: List[str] = Field(default_factory=list, description="Supporting evidence")


# Define custom output schema like above


class VideoRAGSessionConfig(BaseModel):
    """Configuration for VideoRAGSession."""

    vlm_model: str = Field(default="openai/gpt-4o-mini", description="Vision LLM model for litellm")
    similarity_top_k: int = Field(default=5, gt=0, description="Top K for text retrieval")
    image_similarity_top_k: int = Field(default=8, gt=0, description="Top K for image retrieval")
    max_tokens: int = Field(default=1000, gt=0, description="Maximum tokens for VLM response")
    chunk_size: int = Field(default=1000, gt=0, description="Chunk size for text splitting")
    chunk_overlap: int = Field(default=200, ge=0, description="Chunk overlap for text splitting")
    enable_subtitle_check: bool = Field(
        default=True, description="Enable subtitle presence detection"
    )
    subtitle_check_config: Dict[str, Any] = Field(
        default_factory=dict, description="Configuration for subtitle checking"
    )
    max_images_per_query: int = Field(
        default=5,
        ge=1,
        le=20,
        description="Maximum number of images to send to VLM per query (GPT-4V: 10+, Claude: 5+)",
    )

    @field_validator("chunk_overlap")
    @classmethod
    def validate_chunk_overlap(cls, v: int, info) -> int:
        """Validate chunk_overlap is less than chunk_size."""
        chunk_size = info.data.get("chunk_size", 1000)
        if v >= chunk_size:
            raise ValueError(f"chunk_overlap ({v}) must be less than chunk_size ({chunk_size})")
        return v

    @classmethod
    def from_dict(cls, config: Optional[Dict[str, Any]]) -> "VideoRAGSessionConfig":
        """
        Create config from dict (e.g., from video_rag_config.yml).

        Args:
            config: Config dict with optional 'video_rag' key

        Returns:
            VideoRAGSessionConfig instance
        """
        if config is None:
            return cls()

        # Extract video_rag section if present
        rag_config = config.get("video_rag", {}) if isinstance(config, dict) else {}

        return cls.model_validate(rag_config)
