from __future__ import annotations

from typing import Annotated, Literal, Optional, Union

from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator
from pydantic.json_schema import SkipJsonSchema

MAX_REQUEST_ROWS = 200
MAX_CHART_SAMPLE_ROWS = 50
MAX_COLUMNS = 100
MAX_CELL_STRING_LENGTH = 10_000
MAX_ID_LENGTH = 128
MAX_COLUMN_NAME_LENGTH = 256
MAX_QUERY_LENGTH = 8_000
MAX_INSTRUCTIONS_LENGTH = 4_000
MAX_CONTEXT_LENGTH = 8_000
MAX_ORDER_VALUES = 200
MAX_SERIES = 100

BoundedCellString = Annotated[str, Field(max_length=MAX_CELL_STRING_LENGTH)]
JsonScalar = Union[BoundedCellString, float, int, None]
BoundedRow = Annotated[list[JsonScalar], Field(max_length=MAX_COLUMNS)]
BoundedId = Annotated[str, Field(min_length=1, max_length=MAX_ID_LENGTH)]
ColumnName = Annotated[str, Field(min_length=1, max_length=MAX_COLUMN_NAME_LENGTH)]
OrderValue = Annotated[str, Field(max_length=MAX_CELL_STRING_LENGTH)]


class PlanRequest(BaseModel):
    model_config = ConfigDict(extra="forbid")

    columns: list[ColumnName] = Field(default_factory=list, max_length=MAX_COLUMNS)
    data: list[BoundedRow] = Field(default_factory=list, max_length=MAX_REQUEST_ROWS)
    context: Optional[str] = Field(default=None, max_length=MAX_CONTEXT_LENGTH)
    msg_id: BoundedId
    msg_index: int = Field(ge=1, le=100_000)
    conv_id: BoundedId
    artifacts_id: BoundedId
    records_index: int = Field(ge=0, le=100_000)
    artifact_index: int = Field(ge=0, le=100_000)
    query: str = Field(max_length=MAX_QUERY_LENGTH)
    agent_name: Optional[str] = Field(default=None, max_length=256)
    user_instructions: Optional[str] = Field(default=None, max_length=MAX_INSTRUCTIONS_LENGTH)
    previous_chart_plan: Optional[ChartPlan] = None
    persist: bool = True
    row_count: int = Field(default=0, ge=0, le=10_000_000)

    @model_validator(mode="before")
    @classmethod
    def capture_row_count(cls, values):
        if isinstance(values, dict) and isinstance(values.get("data"), list):
            values = {**values, "row_count": len(values["data"])}
        return values

    @field_validator("data", mode="before")
    @classmethod
    def validate_and_limit_data(cls, value):
        if not isinstance(value, list):
            return value
        if len(value) > MAX_REQUEST_ROWS:
            raise ValueError(f"List should have at most {MAX_REQUEST_ROWS} items after validation, not {len(value)}")
        return value[:MAX_CHART_SAMPLE_ROWS]

    @model_validator(mode="after")
    def check_row_widths(self) -> PlanRequest:
        expected_width = len(self.columns)
        if expected_width and any(len(row) != expected_width for row in self.data):
            raise ValueError("each data row must contain exactly one value per column")
        return self


# ---- ChartPlan schema (server-side mirror of FE) ----
class GroupBy(BaseModel):
    model_config = ConfigDict(extra="forbid")
    keyCol: ColumnName
    valueCol: ColumnName
    xIsDate: Optional[bool] = None
    xOrder: Optional[list[OrderValue]] = Field(default=None, max_length=MAX_ORDER_VALUES)
    sortKeys: Optional[list[OrderValue]] = Field(default=None, max_length=MAX_ORDER_VALUES)


class LineGroupByMapping(BaseModel):
    model_config = ConfigDict(extra="forbid")
    x: ColumnName
    groupBy: GroupBy
    duplicateAggregation: SkipJsonSchema[Optional[Literal["first", "sum"]]] = "first"


class LineSeriesColsMapping(BaseModel):
    model_config = ConfigDict(extra="forbid")
    x: ColumnName
    seriesCols: list[ColumnName] = Field(min_length=1, max_length=MAX_SERIES)
    xIsDate: Optional[bool] = None
    xOrder: Optional[list[OrderValue]] = Field(default=None, max_length=MAX_ORDER_VALUES)
    duplicateAggregation: SkipJsonSchema[Optional[Literal["first", "sum"]]] = "first"


class ColumnSeries(BaseModel):
    model_config = ConfigDict(extra="forbid")
    name: str = Field(max_length=256)
    valueCols: list[ColumnName] = Field(min_length=1, max_length=MAX_SERIES)


class LineColumnSeriesMapping(BaseModel):
    """A line/area chart whose categories and series are spread across one data row."""

    model_config = ConfigDict(extra="forbid")
    categoryLabels: list[OrderValue] = Field(min_length=2, max_length=MAX_ORDER_VALUES)
    columnSeries: list[ColumnSeries] = Field(min_length=1, max_length=MAX_SERIES)

    @model_validator(mode="after")
    def check_series_lengths(self) -> LineColumnSeriesMapping:
        category_count = len(self.categoryLabels)
        if any(len(series.valueCols) != category_count for series in self.columnSeries):
            raise ValueError("each columnSeries.valueCols must match categoryLabels length")
        return self


class PieMapping(BaseModel):
    model_config = ConfigDict(extra="forbid")
    labelCol: ColumnName
    valueCol: ColumnName
    aggregate: Optional[Literal["sum", "avg"]] = "sum"


# ---- BAR mapping ----
class BarMapping(BaseModel):
    model_config = ConfigDict(extra="forbid")
    bar: Literal[True]
    x: ColumnName
    groupCol: Optional[ColumnName] = None
    valueCol: ColumnName
    horizontal: Optional[bool] = None
    xOrder: Optional[list[OrderValue]] = Field(default=None, max_length=MAX_ORDER_VALUES)


class BarSeriesColsMapping(BaseModel):
    model_config = ConfigDict(extra="forbid")
    bar: Literal[True]
    x: ColumnName
    seriesCols: list[ColumnName] = Field(min_length=1, max_length=MAX_SERIES)
    horizontal: Optional[bool] = None
    xOrder: Optional[list[OrderValue]] = Field(default=None, max_length=MAX_ORDER_VALUES)


class BarColumnSeriesMapping(BaseModel):
    """A bar chart whose categories and series are spread across one data row."""

    model_config = ConfigDict(extra="forbid")
    bar: Literal[True]
    categoryLabels: list[OrderValue] = Field(min_length=1, max_length=MAX_ORDER_VALUES)
    columnSeries: list[ColumnSeries] = Field(min_length=1, max_length=MAX_SERIES)
    horizontal: Optional[bool] = None

    @model_validator(mode="after")
    def check_series_lengths(self) -> BarColumnSeriesMapping:
        category_count = len(self.categoryLabels)
        if any(len(series.valueCols) != category_count for series in self.columnSeries):
            raise ValueError("each columnSeries.valueCols must match categoryLabels length")
        return self


class WaterfallMapping(BaseModel):
    model_config = ConfigDict(extra="forbid")
    waterfall: Literal[True] = True
    x: ColumnName
    valueCol: ColumnName
    xOrder: Optional[list[OrderValue]] = Field(default=None, max_length=MAX_ORDER_VALUES)


# ---- HEATMAP mapping ----
class HeatmapMapping(BaseModel):
    model_config = ConfigDict(extra="forbid")
    heatmap: Literal[True] = True
    xCol: ColumnName
    yCol: ColumnName
    valueCol: ColumnName
    xOrder: Optional[list[OrderValue]] = Field(default=None, max_length=MAX_ORDER_VALUES)
    yOrder: Optional[list[OrderValue]] = Field(default=None, max_length=MAX_ORDER_VALUES)


class ScatterMapping(BaseModel):
    model_config = ConfigDict(extra="forbid")
    scatter: Literal[True] = True
    xCol: ColumnName
    yCol: ColumnName
    seriesCol: Optional[ColumnName] = None
    xIsDate: Optional[bool] = None


class BubbleMapping(BaseModel):
    model_config = ConfigDict(extra="forbid")
    bubble: Literal[True] = True
    xCol: ColumnName
    yCol: ColumnName
    sizeCol: ColumnName
    seriesCol: Optional[ColumnName] = None
    xIsDate: Optional[bool] = None


Mapping = Union[
    BarMapping,
    BarSeriesColsMapping,
    BarColumnSeriesMapping,
    BubbleMapping,
    HeatmapMapping,
    LineColumnSeriesMapping,
    LineGroupByMapping,
    LineSeriesColsMapping,
    PieMapping,
    ScatterMapping,
    WaterfallMapping,
]


class ChartAppearance(BaseModel):
    model_config = ConfigDict(extra="forbid")
    palette: Literal["default", "brand", "accessible", "monochrome"] = "default"
    showAxisNames: bool = True


class ChartPlan(BaseModel):
    model_config = ConfigDict(extra="forbid")

    kind: Literal["line", "area", "pie", "bar", "waterfall", "heatmap", "scatter", "bubble"]

    title: Optional[str] = Field(default=None, max_length=512)
    donut: Optional[bool] = None
    smooth: Optional[bool] = None
    stacked: Optional[bool] = None
    appearance: SkipJsonSchema[Optional[ChartAppearance]] = None

    mapping: Mapping

    @model_validator(mode="before")
    @classmethod
    def add_legacy_bar_discriminator(cls, values):
        if not isinstance(values, dict) or values.get("kind") != "bar":
            return values
        mapping = values.get("mapping")
        if not isinstance(mapping, dict) or "bar" in mapping:
            return values
        return {**values, "mapping": {**mapping, "bar": True}}

    @model_validator(mode="after")
    def check_mapping_matches_kind(self) -> ChartPlan:
        if self.kind in ("line", "area") and not isinstance(
            self.mapping, (LineColumnSeriesMapping, LineGroupByMapping, LineSeriesColsMapping)
        ):
            raise ValueError("line and area charts require a line mapping")
        if self.kind == "pie" and not isinstance(self.mapping, PieMapping):
            raise ValueError("pie charts require a pie mapping")
        if self.kind == "bar" and not isinstance(
            self.mapping, (BarMapping, BarSeriesColsMapping, BarColumnSeriesMapping)
        ):
            raise ValueError("bar charts require a bar mapping")
        if self.kind == "waterfall" and not isinstance(self.mapping, WaterfallMapping):
            raise ValueError("waterfall charts require a waterfall mapping")
        if self.kind == "heatmap" and not isinstance(self.mapping, HeatmapMapping):
            raise ValueError("heatmap charts require a heatmap mapping")
        if self.kind == "scatter" and not isinstance(self.mapping, ScatterMapping):
            raise ValueError("scatter charts require a scatter mapping")
        if self.kind == "bubble" and not isinstance(self.mapping, BubbleMapping):
            raise ValueError("bubble charts require a bubble mapping")
        return self


class ChartPlanPayload(BaseModel):
    model_config = ConfigDict(extra="forbid")

    chart_type: Literal["none", "line", "area", "pie", "bar", "waterfall", "heatmap", "scatter", "bubble"]
    justification: str = Field(max_length=2_000)
    chart_data: Optional[ChartPlan] = None

    @model_validator(mode="after")
    def check_coherence(self) -> ChartPlanPayload:
        if self.chart_type == "none" and self.chart_data is not None:
            raise ValueError("chart_data must be null when chart_type='none'")
        if self.chart_type != "none" and self.chart_data is None:
            raise ValueError("chart_data is required when chart_type != 'none'")
        if self.chart_data and self.chart_data.kind != self.chart_type:
            # keep kind consistent with outer discriminator
            raise ValueError("chart_type must match chart_data.kind")
        return self


class UpdateChartPlanRequest(BaseModel):
    model_config = ConfigDict(extra="forbid")

    message_id: BoundedId
    chart_plan_index: int = Field(ge=0, le=10_000)
    chart_data: ChartPlan


class ChartPlanDecisionRequest(BaseModel):
    model_config = ConfigDict(extra="forbid")

    message_id: BoundedId
    chart_plan_index: int = Field(ge=0, le=10_000)
    action: Literal["add", "replace"]
    proposal_id: str = Field(min_length=1, max_length=MAX_ID_LENGTH)
    expected_chart_plan: Optional[ChartPlan] = None
    chart_plan: ChartPlanPayload


PlanRequest.model_rebuild()
