"""
Pydantic schemas for cheat sheet structure.

Defines typed data models for all extraction passes and final output.
"""

from dataclasses import dataclass, field
from enum import Enum
from typing import Any, Dict, List, Optional
from pydantic import BaseModel, Field


class ImportanceLevel(str, Enum):
    """Importance level for content prioritization."""
    CRITICAL = "critical"
    HIGH = "high"
    MEDIUM = "medium"
    LOW = "low"


class ConceptItem(BaseModel):
    """A single concept with definition and relationships."""
    name: str = Field(..., description="Concept name")
    definition: str = Field(default="", description="Formal definition")
    intuition: str = Field(default="", description="Plain language explanation")
    assumptions: List[str] = Field(default_factory=list, description="Key assumptions")
    related_concepts: List[str] = Field(default_factory=list, description="Related terms")
    importance: ImportanceLevel = Field(default=ImportanceLevel.MEDIUM)
    source_page: Optional[int] = Field(default=None, description="Source page number")
    
    class Config:
        use_enum_values = True


class FormulaItem(BaseModel):
    """A mathematical formula with context."""
    name: str = Field(..., description="Formula name")
    latex: str = Field(..., description="LaTeX representation")
    variables: Dict[str, str] = Field(default_factory=dict, description="Variable definitions")
    conditions: List[str] = Field(default_factory=list, description="Conditions for use")
    derivation_hint: str = Field(default="", description="Brief derivation note")
    usage_context: str = Field(default="", description="When to use this formula")
    depends_on: List[str] = Field(default_factory=list, description="Prerequisite formulas")
    importance: ImportanceLevel = Field(default=ImportanceLevel.MEDIUM)
    
    class Config:
        use_enum_values = True


class ModelItem(BaseModel):
    """A model or theorem summary."""
    name: str = Field(..., description="Model/theorem name")
    summary: str = Field(..., description="One-paragraph summary")
    objective: str = Field(default="", description="What the model optimizes")
    key_equations: List[str] = Field(default_factory=list, description="Core equations (LaTeX)")
    inputs_outputs: Dict[str, str] = Field(default_factory=dict, description="I/O specification")
    training_procedure: str = Field(default="", description="How it's trained/applied")
    strengths: List[str] = Field(default_factory=list)
    limitations: List[str] = Field(default_factory=list)
    importance: ImportanceLevel = Field(default=ImportanceLevel.MEDIUM)
    
    class Config:
        use_enum_values = True


class AlgorithmItem(BaseModel):
    """An algorithm with pseudocode."""
    name: str = Field(..., description="Algorithm name")
    purpose: str = Field(..., description="What it does")
    pseudocode: str = Field(default="", description="Pseudocode representation")
    inputs: List[str] = Field(default_factory=list, description="Input parameters")
    outputs: List[str] = Field(default_factory=list, description="Output values")
    complexity: str = Field(default="", description="Time/space complexity")
    key_insight: str = Field(default="", description="Main algorithmic insight")
    importance: ImportanceLevel = Field(default=ImportanceLevel.MEDIUM)
    
    class Config:
        use_enum_values = True


class ExampleItem(BaseModel):
    """A worked example or case study."""
    title: str = Field(..., description="Example title")
    description: str = Field(..., description="Problem description")
    solution: str = Field(default="", description="Solution approach")
    key_takeaway: str = Field(default="", description="Main lesson")


class ConnectionItem(BaseModel):
    """A connection to related topics."""
    topic: str = Field(..., description="Related topic name")
    relationship: str = Field(..., description="How it relates")
    external_reference: str = Field(default="", description="Suggested reading")


class CoreConceptsSection(BaseModel):
    """Section containing core concepts."""
    concepts: List[ConceptItem] = Field(default_factory=list)
    summary: str = Field(default="", description="Section summary")


class MathFormulasSection(BaseModel):
    """Section containing mathematical formulas."""
    formulas: List[FormulaItem] = Field(default_factory=list)
    notation_guide: str = Field(default="", description="Notation conventions")


class ModelsTheoremsSection(BaseModel):
    """Section containing models and theorems."""
    models: List[ModelItem] = Field(default_factory=list)
    theorems: List[ModelItem] = Field(default_factory=list)


class AlgorithmSection(BaseModel):
    """Section containing algorithms."""
    algorithms: List[AlgorithmItem] = Field(default_factory=list)


class CheatSheetSchema(BaseModel):
    """
    Complete cheat sheet schema.
    
    This is the final structured output that gets rendered to LaTeX/Markdown.
    """
    title: str = Field(..., description="Cheat sheet title")
    source_document: str = Field(default="", description="Source document name")
    overview: str = Field(default="", description="High-level overview")
    
    # Main sections
    core_concepts: CoreConceptsSection = Field(default_factory=CoreConceptsSection)
    math_formulas: MathFormulasSection = Field(default_factory=MathFormulasSection)
    models_theorems: ModelsTheoremsSection = Field(default_factory=ModelsTheoremsSection)
    algorithms: AlgorithmSection = Field(default_factory=AlgorithmSection)
    
    # Additional sections
    examples: List[ExampleItem] = Field(default_factory=list)
    common_pitfalls: List[str] = Field(default_factory=list)
    connections: List[ConnectionItem] = Field(default_factory=list)
    recap: str = Field(default="", description="One-page recap")
    references: List[str] = Field(default_factory=list)
    
    # Metadata
    page_budget: int = Field(default=12, description="Target page count")
    generation_metadata: Dict[str, Any] = Field(default_factory=dict)
    
    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary for serialization."""
        return self.model_dump()
    
    def get_content_counts(self) -> Dict[str, int]:
        """Get count of items in each section."""
        return {
            "concepts": len(self.core_concepts.concepts),
            "formulas": len(self.math_formulas.formulas),
            "models": len(self.models_theorems.models),
            "theorems": len(self.models_theorems.theorems),
            "algorithms": len(self.algorithms.algorithms),
            "examples": len(self.examples),
            "connections": len(self.connections),
        }


@dataclass
class PassResult:
    """Result from a single summarization pass."""
    success: bool
    pass_name: str
    output: Any = None
    errors: List[str] = field(default_factory=list)
    metrics: Dict[str, Any] = field(default_factory=dict)
    duration_ms: float = 0.0
    tokens_used: int = 0
    
    def to_dict(self) -> Dict[str, Any]:
        return {
            "success": self.success,
            "pass_name": self.pass_name,
            "errors": self.errors,
            "metrics": self.metrics,
            "duration_ms": self.duration_ms,
            "tokens_used": self.tokens_used,
        }
