"""P2HS Calculation IR v1.

Same IR for the local evaluator and the Cloud ``calculation.evaluate``
executor. Authoritative execution uses bounded, pre-resolved dataset snapshots.

Calculation owns compute only. Dataset resolution owns access and scope.
IR nodes must not query records.
"""

from __future__ import annotations

from typing import Annotated, Any, Literal

from pydantic import BaseModel, ConfigDict, Field, model_validator

from src.domain.calculation.errors import CalculationCycleError, CalculationIRError

CALCULATION_IR_VERSION = 1

ValueType = Literal["number", "money", "percent", "date", "boolean", "select", "list", "text"]
ExecutionMode = Literal["local", "authoritative"]
OutputKind = Literal["scalar", "dataset"]
RoundingMode = Literal["round_half_up", "round_half_even"]
DateUnit = Literal["day", "month", "year"]


class TypedInput(BaseModel):
    model_config = ConfigDict(extra="forbid")
    key: str = Field(min_length=1)
    type: ValueType
    required: bool = True
    default: Any = None
    options: list[str] = Field(default_factory=list)
    label: str | None = None

    @model_validator(mode="after")
    def select_needs_options(self) -> TypedInput:
        if self.type == "select" and not self.options:
            raise ValueError(f"select input '{self.key}' requires options")
        return self


class LitExpr(BaseModel):
    model_config = ConfigDict(extra="forbid")
    op: Literal["lit"] = "lit"
    value: Any
    type: ValueType = "number"


class InputExpr(BaseModel):
    model_config = ConfigDict(extra="forbid")
    op: Literal["input"] = "input"
    key: str


class FieldExpr(BaseModel):
    model_config = ConfigDict(extra="forbid")
    op: Literal["field"] = "field"
    key: str


class DatasetExpr(BaseModel):
    """Read a already-resolved Dataset snapshot. Never a query."""

    model_config = ConfigDict(extra="forbid")
    op: Literal["dataset"] = "dataset"
    id: str


class ItemExpr(BaseModel):
    model_config = ConfigDict(extra="forbid")
    op: Literal["item"] = "item"
    name: str = "row"
    field: str | None = None


class UnaryExpr(BaseModel):
    model_config = ConfigDict(extra="forbid")
    op: Literal["neg", "not", "abs"]
    arg: Expr


class BinaryArithExpr(BaseModel):
    model_config = ConfigDict(extra="forbid")
    op: Literal["add", "sub", "mul", "div", "pow"]
    args: list[Expr] = Field(min_length=2, max_length=2)


class CompareExpr(BaseModel):
    model_config = ConfigDict(extra="forbid")
    op: Literal["eq", "neq", "lt", "lte", "gt", "gte"]
    args: list[Expr] = Field(min_length=2, max_length=2)


class IfExpr(BaseModel):
    model_config = ConfigDict(extra="forbid", populate_by_name=True)
    op: Literal["if"] = "if"
    cond: Expr
    then: Expr
    else_: Expr = Field(alias="else")


class VariadicExpr(BaseModel):
    model_config = ConfigDict(extra="forbid")
    op: Literal["min", "max"]
    args: list[Expr] = Field(min_length=1)


class RoundExpr(BaseModel):
    model_config = ConfigDict(extra="forbid")
    op: Literal["round", "floor", "ceil"]
    value: Expr
    places: int = 0
    mode: RoundingMode = "round_half_up"


class ClampExpr(BaseModel):
    model_config = ConfigDict(extra="forbid")
    op: Literal["clamp"] = "clamp"
    value: Expr
    min: Expr
    max: Expr


class DateAddExpr(BaseModel):
    model_config = ConfigDict(extra="forbid")
    op: Literal["date_add"] = "date_add"
    date: Expr
    amount: Expr
    unit: DateUnit = "day"


class DateDiffExpr(BaseModel):
    model_config = ConfigDict(extra="forbid")
    op: Literal["date_diff"] = "date_diff"
    start: Expr
    end: Expr
    unit: DateUnit = "day"


class ListAggExpr(BaseModel):
    model_config = ConfigDict(extra="forbid")
    op: Literal["sum", "count", "avg"]
    over: Expr
    field: str | None = None


class SortExpr(BaseModel):
    """Deterministically order a bounded dataset snapshot by one field."""

    model_config = ConfigDict(extra="forbid")
    op: Literal["sort"] = "sort"
    over: Expr
    field: str = Field(min_length=1)
    direction: Literal["asc", "desc"] = "asc"


class FilterExpr(BaseModel):
    model_config = ConfigDict(extra="forbid")
    op: Literal["filter"] = "filter"
    over: Expr
    item: str = "row"
    where: Expr


class MapExpr(BaseModel):
    model_config = ConfigDict(extra="forbid", populate_by_name=True)
    op: Literal["map"] = "map"
    over: Expr
    item: str = "row"
    as_: Expr = Field(alias="as")


class ListExpr(BaseModel):
    model_config = ConfigDict(extra="forbid")
    op: Literal["list"] = "list"
    items: list[Expr] = Field(default_factory=list)


class RecordExpr(BaseModel):
    model_config = ConfigDict(extra="forbid")
    op: Literal["record"] = "record"
    fields: dict[str, Expr] = Field(default_factory=dict)


Expr = Annotated[
    LitExpr
    | InputExpr
    | FieldExpr
    | DatasetExpr
    | ItemExpr
    | UnaryExpr
    | BinaryArithExpr
    | CompareExpr
    | IfExpr
    | VariadicExpr
    | RoundExpr
    | ClampExpr
    | DateAddExpr
    | DateDiffExpr
    | ListAggExpr
    | SortExpr
    | FilterExpr
    | MapExpr
    | ListExpr
    | RecordExpr,
    Field(discriminator="op"),
]


class DerivedField(BaseModel):
    model_config = ConfigDict(extra="forbid")
    key: str = Field(min_length=1)
    expr: Expr
    type: ValueType | None = None


class OutputSpec(BaseModel):
    model_config = ConfigDict(extra="forbid")
    key: str = Field(min_length=1)
    kind: OutputKind = "scalar"
    from_field: str
    shape: Literal["series", "table", "scalar", "list"] | None = None


class BoundDatasetRef(BaseModel):
    """Named snapshot the evaluator may read. Resolution happens outside IR."""

    model_config = ConfigDict(extra="forbid")
    id: str
    dataset_id: str


class CalculationIR(BaseModel):
    model_config = ConfigDict(extra="forbid")
    schema_version: Literal[1] = 1
    id: str = Field(min_length=1)
    execution: ExecutionMode = "local"
    inputs: list[TypedInput] = Field(default_factory=list)
    datasets: list[BoundDatasetRef] = Field(default_factory=list)
    fields: list[DerivedField] = Field(default_factory=list)
    outputs: list[OutputSpec] = Field(default_factory=list)
    authors: Literal["feature_pack", "codegen"] = "feature_pack"

    @model_validator(mode="after")
    def dag_and_refs(self) -> CalculationIR:
        input_keys = [item.key for item in self.inputs]
        if len(set(input_keys)) != len(input_keys):
            raise CalculationIRError("duplicate input key")
        field_keys = [item.key for item in self.fields]
        if len(set(field_keys)) != len(field_keys):
            raise CalculationIRError("duplicate derived field key")
        dataset_ids = [item.id for item in self.datasets]
        if len(set(dataset_ids)) != len(dataset_ids):
            raise CalculationIRError("duplicate dataset binding id")
        known_fields = set(field_keys)
        known_inputs = set(input_keys)
        known_datasets = set(dataset_ids)
        for field in self.fields:
            _assert_expr_refs(
                field.expr,
                known_inputs=known_inputs,
                known_fields=known_fields,
                known_datasets=known_datasets,
                current=field.key,
            )
        _assert_acyclic(self.fields)
        for output in self.outputs:
            if output.from_field not in known_fields and output.from_field not in known_inputs:
                raise CalculationIRError(
                    f"output '{output.key}' references unknown field '{output.from_field}'"
                )
            if output.kind == "dataset" and output.shape is None:
                raise CalculationIRError(
                    f"dataset output '{output.key}' requires shape"
                )
        if not self.outputs:
            raise CalculationIRError("calculator must declare at least one output")
        return self


def _walk_expr(expr: Expr):
    yield expr
    payload = expr.model_dump(by_alias=True)
    for value in payload.values():
        if isinstance(value, dict) and "op" in value:
            nested = type(expr).model_validate(expr).model_dump()
            break
    else:
        nested = None
    children: list[Any] = []
    for name in (
        "arg",
        "value",
        "cond",
        "then",
        "else_",
        "over",
        "where",
        "as_",
        "date",
        "amount",
        "start",
        "end",
        "min",
        "max",
    ):
        child = getattr(expr, name, None)
        if child is not None and hasattr(child, "op"):
            children.append(child)
    args = getattr(expr, "args", None)
    if args:
        children.extend(args)
    items = getattr(expr, "items", None)
    if items:
        children.extend(items)
    fields = getattr(expr, "fields", None)
    if isinstance(fields, dict):
        children.extend(fields.values())
    for child in children:
        yield from _walk_expr(child)
    _ = nested


def _assert_expr_refs(
    expr: Expr,
    *,
    known_inputs: set[str],
    known_fields: set[str],
    known_datasets: set[str],
    current: str,
) -> None:
    for node in _walk_expr(expr):
        if isinstance(node, InputExpr) and node.key not in known_inputs:
            raise CalculationIRError(f"unknown input '{node.key}'")
        if isinstance(node, FieldExpr):
            if node.key not in known_fields:
                raise CalculationIRError(f"unknown field '{node.key}'")
            if node.key == current:
                raise CalculationCycleError(f"field '{current}' cannot reference itself")
        if isinstance(node, DatasetExpr) and node.id not in known_datasets:
            raise CalculationIRError(
                f"dataset '{node.id}' is not bound; IR cannot query records"
            )


def _field_deps(expr: Expr) -> set[str]:
    return {node.key for node in _walk_expr(expr) if isinstance(node, FieldExpr)}


def _assert_acyclic(fields: list[DerivedField]) -> None:
    deps = {item.key: _field_deps(item.expr) for item in fields}
    visiting: set[str] = set()
    done: set[str] = set()

    def visit(key: str) -> None:
        if key in done:
            return
        if key in visiting:
            raise CalculationCycleError(f"derived field cycle at '{key}'")
        visiting.add(key)
        for dep in deps.get(key, ()):
            if dep in deps:
                visit(dep)
        visiting.remove(key)
        done.add(key)

    for key in deps:
        visit(key)


LitExpr.model_rebuild()
UnaryExpr.model_rebuild()
BinaryArithExpr.model_rebuild()
CompareExpr.model_rebuild()
IfExpr.model_rebuild()
VariadicExpr.model_rebuild()
RoundExpr.model_rebuild()
ClampExpr.model_rebuild()
DateAddExpr.model_rebuild()
DateDiffExpr.model_rebuild()
ListAggExpr.model_rebuild()
SortExpr.model_rebuild()
FilterExpr.model_rebuild()
MapExpr.model_rebuild()
ListExpr.model_rebuild()
RecordExpr.model_rebuild()
DerivedField.model_rebuild()
CalculationIR.model_rebuild()