"""Selection stats schema (stable contract) + validator.

This module is the *external contract* for what `ZenFronts.info(pid)["selection"]`
looks like. The goal is to make downstream code depend on a versioned schema,
not on ad-hoc dict structures.
"""

from __future__ import annotations

from typing import Any, Mapping, TypedDict

import math


SCHEMA_NAME = "zen_fronts.selection_stats"
SCHEMA_VERSION = "1.0.0"


class MetricBlock(TypedDict):
    mean: float
    median: float
    std: float
    q25: float
    q75: float


class SelectionStats(TypedDict):
    # --- schema metadata ---
    schema_name: str
    schema_version: str

    # --- identity / refresh context ---
    point_id: int
    epoch: int
    seed_used: int
    n_samples: int
    percentile: float
    n_active: int
    k: int

    # --- selection outcome ---
    is_winner: bool
    is_loser: bool

    # --- counters aligned with MC sampling ---
    win_count: int
    lose_count: int
    rank_sum: int

    # --- optional advanced stats ---
    quantiles_mode_used: int | None  # 0 exact, 1 streaming (P²)
    quality_score: float | None
    place: MetricBlock | None
    front: MetricBlock | None
    within: MetricBlock | None


def _major(v: str) -> int:
    try:
        return int(v.split(".", 1)[0])
    except Exception as e:  # pragma: no cover
        raise ValueError(f"invalid schema_version: {v!r}") from e


def _as_int(x: Any, *, name: str) -> int:
    if isinstance(x, bool):
        raise TypeError(f"{name} must be int, got bool")
    if isinstance(x, (int,)):
        return int(x)
    # numpy ints etc.
    try:
        return int(x)
    except Exception as e:
        raise TypeError(f"{name} must be int-like, got {type(x).__name__}") from e


def _as_float(x: Any, *, name: str) -> float:
    if isinstance(x, bool):
        raise TypeError(f"{name} must be float, got bool")
    try:
        y = float(x)
    except Exception as e:
        raise TypeError(f"{name} must be float-like, got {type(x).__name__}") from e
    if not math.isfinite(y):
        raise ValueError(f"{name} must be finite, got {y!r}")
    return y


def _as_bool(x: Any, *, name: str) -> bool:
    if isinstance(x, bool):
        return bool(x)
    raise TypeError(f"{name} must be bool, got {type(x).__name__}")


def _parse_metric_block(x: Any, *, name: str) -> MetricBlock:
    if not isinstance(x, Mapping):
        raise TypeError(f"{name} must be a mapping, got {type(x).__name__}")
    mean = _as_float(x.get("mean"), name=f"{name}.mean")
    median = _as_float(x.get("median"), name=f"{name}.median")
    std = _as_float(x.get("std"), name=f"{name}.std")
    q25 = _as_float(x.get("q25"), name=f"{name}.q25")
    q75 = _as_float(x.get("q75"), name=f"{name}.q75")

    if std < 0.0:
        raise ValueError(f"{name}.std must be >= 0")
    if not (q25 <= median <= q75):
        raise ValueError(f"{name} must satisfy q25 <= median <= q75")

    return {"mean": mean, "median": median, "std": std, "q25": q25, "q75": q75}


def validate_selection_stats(obj: Mapping[str, Any]) -> SelectionStats:
    """Validate and normalize a `SelectionStats` object.

    Returns a normalized dict with plain Python ints/floats/bools.

    Compatibility rules
    -------------------
    * schema_name must match exactly.
    * schema_version must have the same MAJOR as this library's SCHEMA_VERSION.
      (Minor/Patch may be newer; downstream can decide how strict it wants to be.)
    """

    if not isinstance(obj, Mapping):
        raise TypeError(f"selection stats must be a mapping, got {type(obj).__name__}")

    schema_name = obj.get("schema_name")
    schema_version = obj.get("schema_version")

    if schema_name != SCHEMA_NAME:
        raise ValueError(f"schema_name mismatch: expected {SCHEMA_NAME!r}, got {schema_name!r}")

    if not isinstance(schema_version, str):
        raise TypeError("schema_version must be str")

    if _major(schema_version) != _major(SCHEMA_VERSION):
        raise ValueError(
            f"incompatible schema_version: got {schema_version!r}, expected MAJOR={_major(SCHEMA_VERSION)}"
        )

    point_id = _as_int(obj.get("point_id"), name="point_id")
    epoch = _as_int(obj.get("epoch"), name="epoch")
    seed_used = _as_int(obj.get("seed_used"), name="seed_used")
    n_samples = _as_int(obj.get("n_samples"), name="n_samples")
    percentile = _as_float(obj.get("percentile"), name="percentile")
    n_active = _as_int(obj.get("n_active"), name="n_active")
    k = _as_int(obj.get("k"), name="k")

    is_winner = _as_bool(obj.get("is_winner"), name="is_winner")
    is_loser = _as_bool(obj.get("is_loser"), name="is_loser")

    win_count = _as_int(obj.get("win_count"), name="win_count")
    lose_count = _as_int(obj.get("lose_count"), name="lose_count")
    rank_sum = _as_int(obj.get("rank_sum"), name="rank_sum")

    quantiles_mode_used_raw = obj.get("quantiles_mode_used")
    quantiles_mode_used: int | None
    if quantiles_mode_used_raw is None:
        quantiles_mode_used = None
    else:
        quantiles_mode_used = _as_int(quantiles_mode_used_raw, name="quantiles_mode_used")
        if quantiles_mode_used not in (0, 1):
            raise ValueError("quantiles_mode_used must be 0, 1, or None")

    quality_score_raw = obj.get("quality_score")
    quality_score = None if quality_score_raw is None else _as_float(quality_score_raw, name="quality_score")

    place_raw = obj.get("place")
    front_raw = obj.get("front")
    within_raw = obj.get("within")

    place = None if place_raw is None else _parse_metric_block(place_raw, name="place")
    front = None if front_raw is None else _parse_metric_block(front_raw, name="front")
    within = None if within_raw is None else _parse_metric_block(within_raw, name="within")

    # Core bounds/invariants
    if n_samples <= 0:
        raise ValueError("n_samples must be > 0")
    if not (0.0 < percentile <= 1.0):
        raise ValueError("percentile must be in (0, 1]")
    if n_active <= 0:
        raise ValueError("n_active must be > 0")
    if not (1 <= k <= n_active):
        raise ValueError("k must be in [1..n_active]")

    if not (0 <= win_count <= n_samples):
        raise ValueError("win_count must be in [0..n_samples]")
    if not (0 <= lose_count <= n_samples):
        raise ValueError("lose_count must be in [0..n_samples]")

    # Rank sum upper bound uses (n_active-1) worst place per sample.
    if not (0 <= rank_sum <= n_samples * max(0, n_active - 1)):
        raise ValueError("rank_sum out of bounds")

    return {
        "schema_name": SCHEMA_NAME,
        "schema_version": str(schema_version),
        "point_id": point_id,
        "epoch": epoch,
        "seed_used": seed_used,
        "n_samples": n_samples,
        "percentile": percentile,
        "n_active": n_active,
        "k": k,
        "is_winner": is_winner,
        "is_loser": is_loser,
        "win_count": win_count,
        "lose_count": lose_count,
        "rank_sum": rank_sum,
        "quantiles_mode_used": quantiles_mode_used,
        "quality_score": quality_score,
        "place": place,
        "front": front,
        "within": within,
    }


# ------------------------------ optional pydantic model ------------------------------

try:  # pragma: no cover
    from pydantic import BaseModel, ConfigDict, Field, field_validator

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

        mean: float
        median: float
        std: float = Field(ge=0.0)
        q25: float
        q75: float

        @field_validator("q75")
        @classmethod
        def _validate_quantiles(cls, v: float, info):  # type: ignore[no-redef]
            data = info.data
            q25 = data.get("q25")
            med = data.get("median")
            if q25 is not None and med is not None and not (q25 <= med <= v):
                raise ValueError("must satisfy q25 <= median <= q75")
            return v


    class SelectionStatsModel(BaseModel):
        """Pydantic parser for downstream (optional dependency)."""

        model_config = ConfigDict(extra="allow")

        schema_name: str = Field(default=SCHEMA_NAME)
        schema_version: str = Field(default=SCHEMA_VERSION)

        point_id: int
        epoch: int
        seed_used: int
        n_samples: int = Field(gt=0)
        percentile: float = Field(gt=0.0, le=1.0)
        n_active: int = Field(gt=0)
        k: int = Field(gt=0)

        is_winner: bool
        is_loser: bool

        win_count: int
        lose_count: int
        rank_sum: int

        quantiles_mode_used: int | None = Field(default=None)
        quality_score: float | None = Field(default=None)
        place: MetricBlockModel | None = None
        front: MetricBlockModel | None = None
        within: MetricBlockModel | None = None

        @field_validator("schema_name")
        @classmethod
        def _schema_name(cls, v: str) -> str:
            if v != SCHEMA_NAME:
                raise ValueError("schema_name mismatch")
            return v

        @field_validator("schema_version")
        @classmethod
        def _schema_version(cls, v: str) -> str:
            if _major(v) != _major(SCHEMA_VERSION):
                raise ValueError("incompatible schema_version")
            return v

        @field_validator("k")
        @classmethod
        def _k_bounds(cls, v: int, info) -> int:
            n_active = info.data.get("n_active")
            if n_active is not None and not (1 <= v <= int(n_active)):
                raise ValueError("k must be in [1..n_active]")
            return v

        @field_validator("win_count")
        @classmethod
        def _win_count(cls, v: int, info) -> int:
            n_samples = info.data.get("n_samples")
            if n_samples is not None and not (0 <= v <= int(n_samples)):
                raise ValueError("win_count out of bounds")
            return v

        @field_validator("lose_count")
        @classmethod
        def _lose_count(cls, v: int, info) -> int:
            n_samples = info.data.get("n_samples")
            if n_samples is not None and not (0 <= v <= int(n_samples)):
                raise ValueError("lose_count out of bounds")
            return v


except Exception:  # pragma: no cover
    # If pydantic isn't installed, keep the module importable.
    MetricBlockModel = None  # type: ignore
    SelectionStatsModel = None  # type: ignore
