"""
Annie SDK Models

DSL models for representing queries. These mirror the same DSL used
by Annie's backend and the pandas-ai-sdk.
"""

from __future__ import annotations

from enum import Enum

from pydantic import BaseModel, Field, model_validator

# =============================================================================
# Enums
# =============================================================================


class Operator(str, Enum):
    """Filter operators."""

    EQUALS = "equals"
    NOT_EQUALS = "not_equals"
    GT = "gt"
    GTE = "gte"
    LT = "lt"
    LTE = "lte"
    CONTAINS = "contains"
    NOT_CONTAINS = "not_contains"
    STARTS_WITH = "starts_with"
    ENDS_WITH = "ends_with"
    IS_NULL = "is_null"
    IS_NOT_NULL = "is_not_null"
    IN = "in"
    NOT_IN = "not_in"
    BETWEEN = "between"


class AggregationFunction(str, Enum):
    """Aggregation functions."""

    COUNT = "count"
    SUM = "sum"
    AVG = "avg"
    MIN = "min"
    MAX = "max"
    COUNT_DISTINCT = "count_distinct"


class SortDirection(str, Enum):
    """Sort direction."""

    ASC = "asc"
    DESC = "desc"


class DateBinningInterval(str, Enum):
    """Date binning intervals."""

    MINUTE = "minute"
    HOUR = "hour"
    DAY = "day"
    WEEK = "week"
    MONTH = "month"
    QUARTER = "quarter"
    YEAR = "year"
    MINUTE_OF_HOUR = "minute_of_hour"
    HOUR_OF_DAY = "hour_of_day"
    DAY_OF_WEEK = "day_of_week"
    DAY_OF_MONTH = "day_of_month"
    DAY_OF_YEAR = "day_of_year"
    WEEK_OF_YEAR = "week_of_year"
    MONTH_OF_YEAR = "month_of_year"
    QUARTER_OF_YEAR = "quarter_of_year"


class NumberBinningStrategy(str, Enum):
    """Number binning strategies."""

    AUTO = "auto"
    FIXED_COUNT = "fixed_count"
    FIXED_WIDTH = "fixed_width"
    CUSTOM = "custom"


class VisualizationType(str, Enum):
    """Visualization types."""

    BAR = "bar"
    PIE = "pie"
    DONUT = "donut"
    LINE = "line"
    COMBO = "combo"
    TABLE = "table"
    KPI = "kpi"
    HORIZONTAL_BAR = "horizontal_bar"


# =============================================================================
# Step Components
# =============================================================================


class FilterCondition(BaseModel):
    """A single filter condition."""

    column: str
    operator: Operator = Operator.EQUALS
    value: str | int | float | bool | list | None = None


class DateBinning(BaseModel):
    """Date binning configuration."""

    interval: DateBinningInterval


class NumberBinning(BaseModel):
    """Number binning configuration."""

    strategy: NumberBinningStrategy = NumberBinningStrategy.AUTO
    bin_count: int | None = None
    bin_width: float | None = None
    boundaries: list[float] | None = None


class AggregationSpec(BaseModel):
    """Aggregation specification."""

    column: str | None = None
    function: AggregationFunction
    label: str | None = None
    format: str | None = None


class BySpec(BaseModel):
    """Group-by specification."""

    column: str
    label: str | None = None
    date_binning: DateBinning | None = None
    number_binning: NumberBinning | None = None


class OrderByColumn(BaseModel):
    """Order-by column specification."""

    column: str
    direction: SortDirection = SortDirection.ASC


# =============================================================================
# Steps
# =============================================================================


class FilterStep(BaseModel):
    """Filter step — applies WHERE conditions."""

    type: str = Field(default="filter", frozen=True)
    conditions: list[FilterCondition] = Field(default_factory=list)


class SummarizeStep(BaseModel):
    """Summarize step — GROUP BY with aggregations."""

    type: str = Field(default="summarize", frozen=True)
    aggregations: list[AggregationSpec] = Field(default_factory=list)
    by: list[BySpec] = Field(default_factory=list)


class OrderByStep(BaseModel):
    """Order-by step — ORDER BY clause."""

    type: str = Field(default="order_by", frozen=True)
    columns: list[OrderByColumn] = Field(default_factory=list)


class LimitStep(BaseModel):
    """Limit step — LIMIT clause."""

    type: str = Field(default="limit", frozen=True)
    value: int = 100


class SelectStep(BaseModel):
    """Select step — SELECT specific columns."""

    type: str = Field(default="select", frozen=True)
    columns: list[str] = Field(default_factory=list)


# =============================================================================
# DataModel
# =============================================================================


class DataModel(BaseModel):
    """A data model representing a query with source and transformation steps."""

    source: str | None = Field(
        None,
        description=(
            "Legacy table name to query. Optional when source_id is provided. "
            "If omitted, the server resolves the table name from source_id."
        ),
    )
    source_id: str | None = Field(
        None,
        description="Source UUID for unambiguous lookup. When provided, takes priority over source for ownership validation.",
    )
    steps: list[FilterStep | SummarizeStep | OrderByStep | LimitStep | SelectStep] = Field(
        default_factory=list, description="Transformation steps"
    )

    @model_validator(mode="after")
    def check_source_or_source_id(self) -> DataModel:
        if not self.source and not self.source_id:
            raise ValueError("At least one of 'source' or 'source_id' must be provided")
        return self
