from __future__ import annotations

from typing import List, Literal, Optional, Union

from pydantic import BaseModel, Field, field_validator, model_validator

JsonScalar = Union[str, float, int, None]
MAX_ROWS = 50 #TODO make it configurable
class PlanRequest(BaseModel):
    columns: List[str]
    data: List[List[JsonScalar]] = Field(default_factory=list)
    context: Optional[str] = None
    msg_id: str
    msg_index: int
    conv_id: str
    artifacts_id: str
    records_index: int
    artifact_index: int
    query: str
    agent_name: Optional[str] = None
    @field_validator("data", mode="before")
    def limit_data(cls, v):
        # avoid huge payloads from FE
        return (v or [])[:MAX_ROWS]

# ---- ChartPlan schema (server-side mirror of FE) ----
class GroupBy(BaseModel):
    keyCol: str
    valueCol: str
    sortKeys: Optional[list[str]]

class LineGroupByMapping(BaseModel):
    x: str
    groupBy: GroupBy

class LineSeriesColsMapping(BaseModel):
    x: str
    seriesCols: List[str]
    xIsDate: Optional[bool] = None
    xOrder: Optional[List[str]] = None

class PieMapping(BaseModel):
    labelCol: str
    valueCol: str
    aggregate: Optional[Literal["sum", "avg"]] = "sum"

# ---- BAR mapping ----
class BarMapping(BaseModel):
    bar: Literal[True] = True
    x: str
    groupCol: str
    valueCol: str
    xOrder: Optional[List[str]] = None

# ---- HEATMAP mapping ----
class HeatmapMapping(BaseModel):
    heatmap: Literal[True] = True
    xCol: str
    yCol: str
    valueCol: str
    xOrder: Optional[List[str]] = None
    yOrder: Optional[List[str]] = None

Mapping = Union[
    LineGroupByMapping,
    LineSeriesColsMapping,
    PieMapping,
    BarMapping,
    HeatmapMapping,
]

class ChartPlan(BaseModel):
    kind: Literal["line", "pie", "bar", "heatmap"]

    title: Optional[str] = None
    donut: Optional[bool] = None
    smooth: Optional[bool] = None
    stacked: Optional[bool] = None

    mapping: Mapping


class ChartPlanPayload(BaseModel):
    chart_type: Literal["none", "line", "pie", "bar", "heatmap"]
    justification: str
    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

