#!/usr/bin/env python3
"""Read-only A/B measurement for the Stage B native-result-handoff rollout.

Compares kanban completion provenance (``kanban_storage.salvage_stats``) across two arms —
flag OFF (baseline) vs flag ON (treatment) — and prints a go/no-go verdict. This tool is
**telemetry only**: it never mutates the board, never touches config, and never runs workers.

Why this exists
---------------
Everything in the grain reframe (Stage A/B) is dormant behind ``[kanban].native_result_handoff``.
The payoff — workers completing via the native tool channel (``handoff_source="tool_submit"``,
counted first-class) instead of relying on salvage/synthesis — can only be proven by running real
kanban workers and reading the provenance. This helper makes that A/B rigorous and repeatable.

Flow
----
  1. python scripts/measure_handoff_ab.py snapshot --out t0.json      # before baseline batch
  2. run a fixed representative batch of N kanban tasks (flag OFF)
  3. python scripts/measure_handoff_ab.py snapshot --out t1.json      # after baseline / before treatment
  4. flip [kanban].native_result_handoff = true; run the same batch shape (flag ON)
  5. python scripts/measure_handoff_ab.py snapshot --out t2.json      # after treatment
  6. python scripts/measure_handoff_ab.py compare --t0 t0.json --t1 t1.json --t2 t2.json

Each arm's metrics are the *delta* in salvage_stats between consecutive snapshots, so they reflect
only that arm's runs. Run on a quiet board with fewer than ``--window`` total terminal runs so the
windowed counts are cumulative and the deltas are exact.

Exit codes (compare): 0 = GO, 1 = NO-GO, 2 = inconclusive / error.
"""
from __future__ import annotations

import argparse
import json
import os
import sys
from datetime import datetime, timezone
from typing import Any

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

# Default A/B window: large enough that, on a quiet board, salvage_stats covers ALL runs so
# snapshot deltas are exact. Capped by salvage_stats itself at 2000.
DEFAULT_WINDOW = 2000


def _now_iso() -> str:
    return datetime.now(tz=timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")


def capture_snapshot(window: int = DEFAULT_WINDOW, *, label: str = "") -> dict[str, Any]:
    """Read salvage_stats now and wrap it with metadata. Read-only."""
    import kanban_storage as ks

    stats = ks.salvage_stats(window=window)
    return {
        "captured_at": _now_iso(),
        "label": label,
        "window": window,
        "stats": stats,
    }


def _stats_of(snapshot: dict[str, Any]) -> dict[str, Any]:
    s = snapshot.get("stats") if isinstance(snapshot, dict) else None
    return s if isinstance(s, dict) else {}


def _int(v: Any) -> int:
    try:
        return int(v or 0)
    except (TypeError, ValueError):
        return 0


def compute_arm(before: dict[str, Any], after: dict[str, Any]) -> dict[str, Any]:
    """Per-source / rate delta between two snapshots for a single arm.

    Uses salvage_stats' own semantics: ``salvaged_runs`` already excludes first-class sources
    (e.g. tool_submit), ``terminal_runs`` counts ALL terminal runs, and ``by_source`` includes
    both salvage and first-class sources. A run with no handoff_source is a clean native/fence
    completion.
    """
    b, a = _stats_of(before), _stats_of(after)
    delta_terminal = _int(a.get("terminal_runs")) - _int(b.get("terminal_runs"))
    delta_salvaged = _int(a.get("salvaged_runs")) - _int(b.get("salvaged_runs"))

    b_src = b.get("by_source") if isinstance(b.get("by_source"), dict) else {}
    a_src = a.get("by_source") if isinstance(a.get("by_source"), dict) else {}
    by_source_delta: dict[str, int] = {}
    for src in set(b_src) | set(a_src):
        d = _int(a_src.get(src)) - _int(b_src.get(src))
        if d:
            by_source_delta[src] = d

    delta_sourced = sum(by_source_delta.values())
    # First-class sourced runs (e.g. tool_submit) = all sourced minus salvage. Derived from
    # storage's own ``salvaged_runs`` (which already excludes first-class sources) rather than a
    # hardcoded source list, so this stays correct if new first-class sources are added. The named
    # breakdown (which first-class source) is visible in ``by_source_delta``.
    delta_first_class = delta_sourced - delta_salvaged
    # Clean (no-source) native/fence completions = terminal minus all sourced runs.
    delta_clean = delta_terminal - delta_sourced

    salvage_rate = (delta_salvaged / delta_terminal) if delta_terminal > 0 else 0.0
    first_class_share = (delta_first_class / delta_terminal) if delta_terminal > 0 else 0.0
    clean_share = (delta_clean / delta_terminal) if delta_terminal > 0 else 0.0
    return {
        "terminal_runs": delta_terminal,
        "salvaged_runs": delta_salvaged,
        "first_class_runs": delta_first_class,
        "clean_runs": delta_clean,
        "salvage_rate": round(salvage_rate, 4),
        "first_class_share": round(first_class_share, 4),
        "clean_share": round(clean_share, 4),
        "by_source_delta": by_source_delta,
    }


def _config_max_salvage_rate(default: float = 0.25) -> float:
    try:
        from core.config import get_kanban_config

        v = get_kanban_config().get("salvage_slo_max", default)
        return float(v) if v is not None else default
    except Exception:
        return default


def evaluate_ab(
    baseline: dict[str, Any],
    treatment: dict[str, Any],
    *,
    min_sample: int = 20,
    max_salvage_rate: float | None = None,
) -> dict[str, Any]:
    """Compare two arm metrics (from :func:`compute_arm`) and produce a go/no-go verdict.

    Returns a dict with ``verdict`` in {"go", "no-go", "inconclusive"}, ``reasons`` (list of
    short strings), and a ``warnings`` list (non-blocking observations).
    """
    reasons: list[str] = []
    warnings: list[str] = []

    b_n = _int(baseline.get("terminal_runs"))
    t_n = _int(treatment.get("terminal_runs"))

    # Inconclusive: not enough runs in either arm to draw a conclusion.
    if b_n <= 0 or t_n <= 0:
        reasons.append(
            f"non-positive terminal runs (baseline={b_n}, treatment={t_n}); "
            "check snapshot order and that batches actually ran"
        )
        return {"verdict": "inconclusive", "reasons": reasons, "warnings": warnings}
    if b_n < min_sample or t_n < min_sample:
        reasons.append(
            f"insufficient sample (baseline={b_n}, treatment={t_n}; need >= {min_sample} per arm)"
        )
        return {"verdict": "inconclusive", "reasons": reasons, "warnings": warnings}

    b_rate = float(baseline.get("salvage_rate") or 0.0)
    t_rate = float(treatment.get("salvage_rate") or 0.0)
    t_first_class = _int(treatment.get("first_class_runs"))

    go = True
    if t_first_class <= 0:
        go = False
        reasons.append(
            "treatment produced ZERO first-class (tool_submit) completions — the native channel "
            "was not used (teaching/availability/plumbing gap); flag had no effect"
        )
    else:
        reasons.append(
            f"treatment used the native tool channel ({t_first_class} first-class completions)"
        )

    if t_rate < b_rate:
        reasons.append(f"salvage_rate dropped {b_rate:.0%} -> {t_rate:.0%}")
    elif t_rate == b_rate:
        go = False
        reasons.append(f"salvage_rate unchanged at {t_rate:.0%}")
    else:
        go = False
        reasons.append(f"salvage_rate INCREASED {b_rate:.0%} -> {t_rate:.0%} (regression)")

    eff_max = max_salvage_rate if max_salvage_rate is not None else _config_max_salvage_rate()
    if t_rate > eff_max:
        warnings.append(
            f"treatment salvage_rate {t_rate:.0%} still exceeds SLO max {eff_max:.0%} "
            "(relative improvement noted, but not yet under SLO)"
        )

    warnings.append(
        "manual guardrail: confirm gave_up / parse_failure outcomes did NOT increase "
        "(salvage_stats does not expose outcome counts) — check `mugi-kanban health`"
    )
    return {"verdict": "go" if go else "no-go", "reasons": reasons, "warnings": warnings}


def _load(path: str) -> dict[str, Any]:
    with open(path, encoding="utf-8") as fh:
        return json.load(fh)


def _cmd_snapshot(args: argparse.Namespace) -> int:
    snap = capture_snapshot(window=int(args.window), label=args.label or "")
    payload = json.dumps(snap, indent=2, ensure_ascii=False)
    if args.out:
        with open(args.out, "w", encoding="utf-8") as fh:
            fh.write(payload + "\n")
        print(f"snapshot written: {args.out}")
        s = snap["stats"]
        print(
            f"  terminal_runs={s.get('terminal_runs')} salvaged_runs={s.get('salvaged_runs')} "
            f"salvage_rate={float(s.get('salvage_rate') or 0):.0%} by_source={s.get('by_source')}"
        )
    else:
        print(payload)
    return 0


def _print_arm(name: str, arm: dict[str, Any]) -> None:
    print(
        f"  {name}: terminal={arm['terminal_runs']} salvaged={arm['salvaged_runs']} "
        f"first_class={arm['first_class_runs']} clean={arm['clean_runs']} "
        f"salvage_rate={arm['salvage_rate']:.0%} first_class_share={arm['first_class_share']:.0%}"
    )
    if arm["by_source_delta"]:
        print(f"      by_source_delta={arm['by_source_delta']}")


def _cmd_compare(args: argparse.Namespace) -> int:
    try:
        t0, t1, t2 = _load(args.t0), _load(args.t1), _load(args.t2)
    except (OSError, json.JSONDecodeError) as exc:
        print(f"error loading snapshots: {exc}", file=sys.stderr)
        return 2

    baseline = compute_arm(t0, t1)
    treatment = compute_arm(t1, t2)
    verdict = evaluate_ab(
        baseline,
        treatment,
        min_sample=int(args.min_sample),
        max_salvage_rate=args.max_salvage_rate,
    )

    print("Stage B native-result-handoff A/B")
    _print_arm("baseline (flag off)", baseline)
    _print_arm("treatment (flag on)", treatment)
    print(f"VERDICT: {verdict['verdict'].upper()}")
    for r in verdict["reasons"]:
        print(f"  - {r}")
    for w in verdict["warnings"]:
        print(f"  ! {w}")

    if args.json_out:
        with open(args.json_out, "w", encoding="utf-8") as fh:
            json.dump(
                {"baseline": baseline, "treatment": treatment, "verdict": verdict},
                fh,
                indent=2,
            )

    return {"go": 0, "no-go": 1, "inconclusive": 2}[verdict["verdict"]]


def build_arg_parser() -> argparse.ArgumentParser:
    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    sub = p.add_subparsers(dest="cmd", required=True)

    sp = sub.add_parser("snapshot", help="capture salvage_stats now (read-only)")
    sp.add_argument("--out", default="", help="write snapshot JSON to this path (else stdout)")
    sp.add_argument("--window", type=int, default=DEFAULT_WINDOW, help="salvage_stats window")
    sp.add_argument("--label", default="", help="optional label stored in the snapshot")
    sp.set_defaults(func=_cmd_snapshot)

    cp = sub.add_parser("compare", help="compare baseline (t0->t1) vs treatment (t1->t2)")
    cp.add_argument("--t0", required=True, help="snapshot before the baseline batch")
    cp.add_argument("--t1", required=True, help="snapshot after baseline / before treatment")
    cp.add_argument("--t2", required=True, help="snapshot after the treatment batch")
    cp.add_argument("--min-sample", type=int, default=20, help="min terminal runs per arm")
    cp.add_argument("--max-salvage-rate", type=float, default=None, help="SLO max (default: config)")
    cp.add_argument("--json-out", default="", help="optional path to write the full comparison JSON")
    cp.set_defaults(func=_cmd_compare)
    return p


def main(argv: list[str] | None = None) -> int:
    args = build_arg_parser().parse_args(argv)
    return args.func(args)


if __name__ == "__main__":
    raise SystemExit(main())
