"""Assemble the experiment report payload from persisted ``_ab_results`` rows.

One JSON-serializable payload per **experiment** — the shared contract the
offline HTML readout (WP3) and the explore cockpit (WP6/WP7) both consume
(data-contract-and-reporting.md §5.3; m3-implementation-plan.md D6). Kept in
documented lockstep with the renderer-side ``web/src/shared/payload.ts``
(WP3) — same keys, same units.

The payload derives from **stored** rows (including their persisted
``method_params``/``alpha``/``srm_flag``), not from re-evaluating what the
current YAML would produce — the report shows *what actually ran* (the donor
doctrine). Header metadata comes from the experiment config (the truth);
``_ab_experiments`` is informational only and never read here.

Units and null discipline (§5.3): timestamps are integer ms-epoch (UTC);
every nullable numeric passes through :func:`_num_or_none`, so NaN **and
±inf** become JSON ``null`` (H5 zero-denominator NaNs; the verdict's
``pair_mde`` ``math.inf`` = "configured but unavailable" — the rationale
strings carry the explanation). No NaN, numpy scalar, or datetime object ever
reaches the returned dict.

Cost discipline: per-point MDE reads the **stored** columns only — the
read-time D5(b) solve runs once per pair on the latest cutoff (the verdict),
never per historical point (a statsmodels root-solve × O(cutoffs) would cost
seconds-to-minutes at sub-day cadence; review finding).
"""

from __future__ import annotations

import math
from collections.abc import Mapping
from datetime import datetime
from itertools import combinations
from pathlib import Path
from typing import Any

import numpy as np

from abkit.config.experiment_config import ComparisonConfig, ExperimentConfig
from abkit.config.metric_config import MetricConfig
from abkit.config.project_config import ProjectConfig
from abkit.database.internal_tables import InternalTablesManager
from abkit.database.manager import BaseDatabaseManager
from abkit.loaders.exposure_source import (
    EmptyCohortError,
    render_assignment_sql,
    validate_and_snapshot,
)
from abkit.pipeline.readout import ExperimentReadout, PairVerdict, evaluate, srm_summary
from abkit.reporting.calibration import build_calibration_block
from abkit.utils.datetime_utils import to_naive_utc
from abkit.utils.json_utils import json_loads

#: the payload schema version (§5.3) — bump on breaking key/unit changes
PAYLOAD_VERSION = 1

#: total point budget across metrics × pairs × cutoffs; exceeding it clips
#: every series to its trailing window and adds a top-level payload warning
#: (the donor capped at 1500 points for one series; an experiment payload
#: carries many series, so the budget is global)
REPORT_POINT_BUDGET = 20_000


def _ms(value: Any) -> int:
    """Coerce a datetime / datetime64 to integer ms-epoch (UTC)."""
    if isinstance(value, datetime):
        value = to_naive_utc(value)
    return int(np.datetime64(value, "ms").astype("int64"))


def _num_or_none(value: Any) -> float | None:
    """Pass a stored Nullable number through; NaN/±inf/None become ``None``.

    The inf clause extends the donor helper: ``pair_mde`` uses ``math.inf``
    for "``calculate_mde: true`` but columns are NULL", and JSON ``Infinity``
    is invalid in browsers — serialization only, no statistical number
    changes.
    """
    if value is None:
        return None
    try:
        f = float(value)
    except (TypeError, ValueError):
        return None
    return f if math.isfinite(f) else None


def _flag01(value: Any) -> int:
    """A stored Bool cell (bool / 0-1 int / numpy) as a terse 0/1 point flag."""
    return 1 if value else 0


def _reject_flag(value: Any) -> int | None:
    """``reject`` is Nullable(Bool): None = inference withheld (demoted)."""
    if value is None:
        return None
    return 1 if value else 0


def _parse_json_cell(value: Any) -> Any:
    """Parse a stored canonical-JSON cell (str | None); never raises."""
    if value is None or value == "":
        return None
    if isinstance(value, (dict, list)):
        return value
    try:
        return json_loads(value)
    except (TypeError, ValueError):
        return None


def _stored_pair_mde(row: dict) -> float | None:
    """The per-point pair MDE from the **stored** columns only (D5(b) combine).

    Deliberately no read-time ``pair_mde`` solve here: that runs a statsmodels
    root-solve per t/z row, and cumulative windows give every cutoff a unique
    ``(size, ratio)`` — so the whole series would cost O(cutoffs) solves
    (seconds-to-minutes at sub-day cadence; review finding). The point series
    honestly shows *what actually ran*: MDE where the row computed it
    (``calculate_mde: true``), null otherwise. The read-time D5(b) fallback
    stays where the FLAT decision needs it — the **verdict**, one solve per
    pair on the latest cutoff (readout ``pair_mde``).

    The **both-present guard** mirrors readout ``pair_mde``: enrich NULLs a
    non-finite solve, so a half-present pair means one arm is undetectable —
    taking the finite arm alone would fake adequate power (the exact trap the
    verdict path guards; review finding). Half-present ⇒ null, never the
    finite arm."""
    mde_1, mde_2 = _num_or_none(row.get("mde_1")), _num_or_none(row.get("mde_2"))
    if mde_1 is None or mde_2 is None:
        return None
    return max(abs(mde_1), abs(mde_2))


def _point(row: dict) -> dict:
    """One ``_ab_results`` row → the terse §5.3 series point.

    ``mde`` is the per-point pair MDE from the stored columns (:func:
    `_stored_pair_mde`). Demoted rows pass their NULL test columns through as
    nulls; sizes stay real. ``v1/v2/sd1/sd2/cv1/cv2`` are the per-arm stored
    values/stds/CUPED covariate means — the WP3 additive extension (no v-bump)
    feeding the §5.2 "variant means/lift" and §3 view-2 renderings; ``cv*`` is
    null for non-CUPED rows.
    """
    return {
        "t": _ms(row["end_ts"]),
        "ed": _num_or_none(row.get("elapsed_days")),
        "e": _num_or_none(row.get("effect")),
        "lo": _num_or_none(row.get("left_bound")),
        "hi": _num_or_none(row.get("right_bound")),
        "p": _num_or_none(row.get("pvalue")),
        "rj": _reject_flag(row.get("reject")),
        "s1": int(row.get("size_1") or 0),
        "s2": int(row.get("size_2") or 0),
        "v1": _num_or_none(row.get("value_1")),
        "v2": _num_or_none(row.get("value_2")),
        "sd1": _num_or_none(row.get("std_1")),
        "sd2": _num_or_none(row.get("std_2")),
        "cv1": _num_or_none(row.get("cov_value_1")),
        "cv2": _num_or_none(row.get("cov_value_2")),
        "mde": _stored_pair_mde(row),
        "hz": _flag01(row.get("is_horizon")),
        "blk": _flag01(row.get("decision_blocked")),
        "ins": _flag01(row.get("insufficient_data")),
    }


def _verdict_to_payload(verdict: PairVerdict) -> dict:
    """A WP1 :class:`PairVerdict` → the JSON-safe §5.3 verdict entry."""
    return {
        "metric": verdict.metric,
        "pair": {"c": verdict.name_1, "t": verdict.name_2},
        "verdict": verdict.verdict,
        "rationale": list(verdict.rationale),
        "caveats": list(verdict.caveats),
        "significant": bool(verdict.significant),
        "effect": _num_or_none(verdict.effect),
        "pvalue": _num_or_none(verdict.pvalue),
        "lo": _num_or_none(verdict.left_bound),
        "hi": _num_or_none(verdict.right_bound),
        "alpha": _num_or_none(verdict.alpha),
        "mde": _num_or_none(verdict.mde),
        "min_effect": _num_or_none(verdict.min_effect),
        "end_ts": _ms(verdict.end_ts) if verdict.end_ts is not None else None,
        "elapsed_days": _num_or_none(verdict.elapsed_days),
        "is_horizon": bool(verdict.is_horizon),
        # §6.5 representativeness chip: the weekly-cycle coverage fraction on an
        # early decisive verdict (null otherwise); the report promotes it to a
        # rendered chip rather than leaving it a caveat bullet.
        "weekly_cycle_pct": _num_or_none(verdict.weekly_cycle_pct),
        "guardrails": [
            {
                "metric": g.metric,
                "pair": {"c": g.name_1, "t": g.name_2},
                "regressed": bool(g.regressed),
                "effect": _num_or_none(g.effect),
                "desired_direction": g.desired_direction,
            }
            for g in verdict.guardrails
        ],
    }


def _window_filter(rows: list[dict], start: datetime | None, end: datetime | None) -> list[dict]:
    """Keep rows with ``start <= end_ts <= end`` (both bounds inclusive).

    Pinning ``end`` before the latest cutoff replays the report as-of that
    moment — the readout then verdicts on the historical state (the donor's
    replay stance).
    """
    if start is None and end is None:
        return rows
    kept = []
    for row in rows:
        end_ts = row["end_ts"]
        if start is not None and end_ts < start:
            continue
        if end is not None and end_ts > end:
            continue
        kept.append(row)
    return kept


def _metric_entry(
    comparison: ComparisonConfig,
    rows: list[dict],
    variants: list[str],
    metric_configs: Mapping[str, MetricConfig] | None,
) -> dict:
    """One configured comparison → the §5.3 ``metrics[]`` entry.

    ``description`` comes from the metric YAML config (D6 — ``_ab_experiments``
    stores only the experiment description). ``method``/``query``/``alpha``
    prefer the latest stored row (what actually ran); config values are the
    fallback for a never-run comparison. Provenance is deduped to one
    ``query`` entry per metric; the rendered SQL never enters the payload.
    """
    metric_config = (metric_configs or {}).get(comparison.metric)

    by_pair: dict[tuple[str, str], list[dict]] = {}
    for row in rows:
        by_pair.setdefault((str(row["name_1"]), str(row["name_2"])), []).append(row)

    latest: dict | None = None
    for group in by_pair.values():
        candidate = group[-1]
        if latest is None or candidate["end_ts"] > latest["end_ts"]:
            latest = candidate

    method_params = None
    if latest is not None:
        method_params = _parse_json_cell(latest.get("method_params"))
    if method_params is None:
        method_params = _parse_json_cell(comparison.method.canonical_params_json) or {}

    warnings: list[str] = []
    seen_warnings: set[str] = set()
    for row in rows:
        for warning in _parse_json_cell(row.get("warnings")) or []:
            text = str(warning)
            if text not in seen_warnings:
                seen_warnings.add(text)
                warnings.append(text)

    pairs = []
    for name_1, name_2 in combinations(variants, 2):
        group = by_pair.get((name_1, name_2), [])
        pairs.append(
            {
                "c": name_1,
                "t": name_2,
                "series": [_point(row) for row in group],
                "diag": _parse_json_cell(group[-1].get("diagnostics")) if group else None,
            }
        )

    return {
        "name": comparison.metric,
        "description": metric_config.description if metric_config is not None else None,
        "main": bool(comparison.is_main_metric),
        "guardrail": bool(comparison.is_guardrail),
        "method": {
            "name": comparison.method.name,
            "params": method_params,
            "id": comparison.method.method_config_id,
            "alpha": _num_or_none(latest.get("alpha")) if latest is not None else None,
        },
        "query": str(latest.get("metric_query")) if latest is not None else None,
        "pairs": pairs,
        "warnings": warnings,
    }


def _clip_to_budget(metrics: list[dict], max_points: int) -> str | None:
    """Enforce the global point budget by trailing-window truncation.

    Runs AFTER the readout evaluated the full series — clipping affects only
    what gets baked, never the verdict. Returns the payload warning when
    anything was dropped (no silent caps)."""
    series_list = [pair["series"] for metric in metrics for pair in metric["pairs"]]
    non_empty = [s for s in series_list if s]
    total = sum(len(s) for s in non_empty)
    if total <= max_points or not non_empty:
        return None
    allowed = max(1, max_points // len(non_empty))
    dropped = 0
    for metric in metrics:
        for pair in metric["pairs"]:
            series = pair["series"]
            if len(series) > allowed:
                dropped += len(series) - allowed
                pair["series"] = series[-allowed:]
    return (
        f"payload clipped to the trailing {allowed} cutoffs per series "
        f"({total} points exceeded the {max_points}-point budget; "
        f"{dropped} older points dropped — pin start/end to window explicitly)"
    )


def build_report_payload(
    experiment: ExperimentConfig,
    tables: InternalTablesManager,
    *,
    project: ProjectConfig | None = None,
    metric_configs: Mapping[str, MetricConfig] | None = None,
    generated_at: str | None = None,
    start: datetime | None = None,
    end: datetime | None = None,
    max_points: int = REPORT_POINT_BUDGET,
    manager: BaseDatabaseManager | None = None,
    project_root: Path | None = None,
    cohort_counts: dict[str, int] | None = None,
) -> dict:
    """Read ``_ab_results`` for one experiment and assemble the §5.3 payload.

    ``start``/``end`` bound the report window on ``end_ts`` (both inclusive);
    pinning ``end`` replays the readout as-of a historical cutoff.
    ``generated_at`` is a caller-supplied preformatted string — the builder
    never reads the wall clock (determinism; the CLI owns formatting).
    ``project`` resolves the effective correction for the readout (BH is
    read-time) and names the payload; ``metric_configs`` supplies the metric
    descriptions (D6). On a never-run project (no ``_ab_results`` table) every
    key is still present with empty series — reporting never creates schema.

    The SRM chip's observed counts follow the m8 WP4 source-mode switch:
    ``cohort_counts`` (a snapshot the caller already validated — the
    explore/run/validate call sites thread theirs through, so one invocation
    never executes the source twice) wins outright; otherwise copy mode reads
    the persisted ``_ab_exposures`` and direct mode — given ``manager`` +
    ``project_root`` — snapshots the live assignment source (an empty source
    reads as zero counts, like a missing table did). A direct-mode call
    without a ``manager`` shows ZEROS: the persisted table is not maintained
    in direct mode, so reading it could silently present a stale copy-era
    cohort as current — zeros are the honest "no source available" face.
    """
    if start is not None:
        start = to_naive_utc(start)
    if end is not None:
        end = to_naive_utc(end)

    grid = experiment.grid()
    variants = list(experiment.assignment.variants)

    declared_pairs = set(combinations(variants, 2))
    ready = tables.results_table_exists()
    declared_by_comparison: dict[str, list[dict]] = {}
    rows_by_comparison: dict[str, list[dict]] = {}
    stale_by_metric: dict[str, int] = {}
    for comparison in experiment.comparisons:
        loaded = (
            tables.load_results(
                experiment.name, comparison.metric, comparison.method.method_config_id
            )
            if ready
            else []
        )
        # pairs outside the declared variants (a mid-flight arm rename) are
        # not chartable; drop them from EVERY payload surface — series, look,
        # period, latest-row method block — loudly, never silently
        declared = [
            row for row in loaded if (str(row["name_1"]), str(row["name_2"])) in declared_pairs
        ]
        if len(declared) != len(loaded):
            stale_by_metric[comparison.metric] = len(loaded) - len(declared)
        declared_by_comparison[comparison.metric] = declared
        # the chart + verdict are as-of the window; the SRM block is not
        rows_by_comparison[comparison.metric] = _window_filter(declared, start, end)

    charted_rows = [row for rows in rows_by_comparison.values() for row in rows]
    declared_rows = [row for rows in declared_by_comparison.values() for row in rows]
    readout: ExperimentReadout = evaluate(experiment, charted_rows, project=project)

    # SRM block — CURRENT experiment health, window-independent (§6 "SRM loud"):
    # flag/pvalue come from the latest persisted row overall (not the latest
    # *charted* row), so a pinned/empty replay never silences a failing gate.
    # `observed` is the whole-cohort current split; at daily+ (χ²) it is the
    # gate's own input, while below 1d (the anytime e-process, WP5) the flag is
    # the latest look's running verdict — so a since-rebalanced sub-day cohort
    # can read "FAILED" with a now-even split (an imbalance WAS detected). The
    # chip names the gate via `srm.kind` so this reads correctly, not as χ².
    srm_flag, srm_pvalue = srm_summary(experiment, declared_rows)
    observed_counts: dict[str, int] = {}
    if cohort_counts is not None:
        observed_counts = dict(cohort_counts)
    elif ready:
        if experiment.assignment.cohort_copy.enabled:
            if tables.exposures_table_exists():
                observed_counts = tables.get_exposure_counts(experiment.name)
        elif manager is not None:
            # direct mode (m8 WP4): the persisted copy is not written — snapshot
            # the live source for the chip's counts. Empty source → zeros (the
            # same face a missing table showed); genuine contract violations
            # propagate to the caller's best-effort report guard. Without a
            # manager the counts stay ZERO — never the possibly-stale
            # persisted table (see the docstring).
            try:
                rendered = render_assignment_sql(manager, experiment, project_root, grid)
                observed_counts = validate_and_snapshot(manager, experiment, rendered).counts
            except EmptyCohortError:
                observed_counts = {}
    # zero-fill declared variants, mirroring the driver's SRM gate input
    observed = {variant: int(observed_counts.get(variant, 0)) for variant in variants}

    metrics = [
        _metric_entry(comparison, rows_by_comparison[comparison.metric], variants, metric_configs)
        for comparison in experiment.comparisons
    ]

    # calibration block (M4): the latest `abk validate` matrix for this experiment,
    # read-only (never creates schema — the `aa_runs_table_exists` guard). Independent
    # of `_ab_results` (a validate-only project still shows its matrix). `null` until
    # the first validate run, so the chip reads the M3 "uncalibrated" empty state
    # (data-contract §5.3; the M4 shape lands without a payload v-bump).
    calibration = None
    if tables.aa_runs_table_exists():
        calibration = build_calibration_block(tables.get_aa_runs(experiment.name))

    warnings = list(readout.warnings)
    for metric_name, dropped in stale_by_metric.items():
        warnings.append(
            f"{experiment.name}/{metric_name}: {dropped} persisted rows are for "
            "variant pairs outside the declared variants (renamed arms?) — not "
            "charted"
        )
    if ready:
        # the driver's orphan scan, surfaced on the read path too: an edited
        # method leaves the old series in _ab_results and the report would
        # otherwise show a silently truncated history
        stored_ids = tables.list_method_config_ids(experiment.name)
        for comparison in experiment.comparisons:
            orphaned = {
                mc_id for (metric_name, mc_id) in stored_ids if metric_name == comparison.metric
            } - {comparison.method.method_config_id}
            if orphaned:
                warnings.append(
                    f"{experiment.name}/{comparison.metric}: {len(orphaned)} orphaned "
                    "method_config_id series in _ab_results (the BI chart will "
                    "show duplicate stabilization lines) — run `abk clean`"
                )
    clip_warning = _clip_to_budget(metrics, max_points)
    if clip_warning is not None:
        warnings.append(clip_warning)

    # look counter (§4): n = charted cutoffs where inference actually happened
    # (at least one non-demoted row); planned = the one-enumeration grid length
    informative_cutoffs = {
        row["end_ts"] for row in charted_rows if not row.get("insufficient_data")
    }
    latest_end_ts = max((row["end_ts"] for row in charted_rows), default=None)

    return {
        "v": PAYLOAD_VERSION,
        "experiment": experiment.name,
        "project": project.name if project is not None else None,
        "generated_at": generated_at,
        "description": experiment.description,
        "period": {
            "start": _ms(grid.start_ts),
            # 0 = no persisted cutoffs (the donor's empty sentinel); start
            # and horizon stay real — they are config facts, not data facts
            "end": _ms(latest_end_ts) if latest_end_ts is not None else 0,
            "horizon": _ms(grid.horizon_ts),
        },
        "cadence_seconds": experiment.cadence_seconds_min(),
        "tz": experiment.timezone,
        "arms": variants,
        "srm": {
            "flag": bool(srm_flag),
            "pvalue": _num_or_none(srm_pvalue),
            "observed": observed,
            "expected": {
                variant: float(split)
                for variant, split in experiment.assignment.expected_split.items()
            },
            # which gate produced flag/pvalue, so the chip names the right test:
            # χ² at daily+, the anytime-valid e-process below 1d (WP5). Not
            # persisted, so derived from the current cadence (statistics-changes §4.2).
            "kind": "sequential_multinomial" if experiment.is_sub_day() else "chi2",
        },
        # the latest A/A validate matrix (M4); null until first `abk validate`.
        # The M4 shape (fpr, peeking_fpr, headline, matrix_rows, report_link) fills
        # in without a payload v-bump (§5.3; every field renderer-side optional).
        "calibration": calibration,
        "verdicts": [_verdict_to_payload(v) for v in readout.verdicts],
        "metrics": metrics,
        "look": {"n": len(informative_cutoffs), "planned": len(grid)},
        # injected by the explore server at serve time; null in a baked report
        "endpoints": {
            "save_url": None,
            "recompute_url": None,
            "reload_url": None,
            "validate_url": None,
        },
        "warnings": warnings,
    }
