"""
Pydantic models for structured outputs from the LLM.
Ensures clean JSON that Python can reliably execute without crashing.
"""

from enum import Enum
from typing import Any, Literal, Optional

from pydantic import BaseModel, Field


class RouteQuery(BaseModel):
    """Structured routing decision for multi-source RAG."""

    datasource: Literal["sql", "vector_store"] = Field(
        description="Route to 'sql' for structured/numerical data (sales, counts, aggregations). "
        "Route to 'vector_store' for policies, documentation, procedures, FAQs."
    )
    reason: str = Field(description="Brief explanation of the routing decision.")


class ChartType(str, Enum):
    """Supported visualization types for data analysis."""

    LINE = "line"  # Trend over time
    BAR = "bar"  # Categorical comparisons
    BARH = "barh"  # Horizontal bar
    PIE = "pie"  # Part-to-whole
    SCATTER = "scatter"  # Correlation
    HISTOGRAM = "histogram"  # Distribution
    BOX = "box"  # Statistical distribution
    HEATMAP = "heatmap"  # Matrix/correlation


class SQLQueryOutput(BaseModel):
    """Structured output for SQL query generation."""

    query: str = Field(description="The SQL query to execute. Must be valid SQLite syntax.")
    explanation: str = Field(
        description="Brief explanation of what the query does and why it answers the question."
    )


class VisualizationConfig(BaseModel):
    """Structured output for chart configuration."""

    chart_type: ChartType = Field(
        description="The type of chart that best represents the data."
    )
    title: str = Field(description="Descriptive title for the chart.")
    x_column: Optional[str] = Field(
        default=None, description="Column name for x-axis (or categories for bar/pie)."
    )
    y_column: Optional[str] = Field(
        default=None, description="Column name for y-axis (or values for bar/pie)."
    )
    hue_column: Optional[str] = Field(
        default=None, description="Column for color grouping (e.g., in bar charts)."
    )