"""Standalone checkpoint eval: fresh process, single GPU, golden-proven generator.

Promotion of the .tmp/eval_gate.py driver (task #26) into a parameterized
entry point. The design facts it carries forward are all incident-derived:

  * FRESH PROCESS, never in-trainer: the train->eval handoff inside one
    process produced the rank race (peers built a second 7B while rank 0
    saved) and the orphan-GPU-process incident. A checkpoint on disk plus a
    fresh process has neither failure mode.
  * EVERYTHING under the __main__ guard: the verifier's scoring pool uses
    the 'spawn' context, which RE-IMPORTS __main__ in every worker -- the
    unguarded first driver had all 8 workers load their own 7B onto the
    eval GPU.
  * greedy generation via `golden_generate_fn` (in-repo Engine, same encode
    as the golden regression: BOS default, truncation 768). The anchors and
    banked SFT rows were measured with HF generate; the generator swap is
    certified by the golden regression + checkpoint roundtrip gate, and the
    generator is recorded in the report config.
  * Atomic report write (.partial -> os.replace): a half-written report
    masquerades as a complete one.

Phase-0 usage (the Arm B result side of the primary McNemar):

    CUDA_VISIBLE_DEVICES=0 python scripts/eval_checkpoint.py \
        --checkpoint runs/dpo_arm_b_response/checkpoints/micro_XXXXXXXX \
        --name arm-b-final --gsm8k-n 1319 --math500-n 0

`--quality-gate` additionally runs the SFT-stage quality gate + ledger
records; leave it off for non-SFT checkpoints (the gate's anchors are
SFT-specific).
"""

import argparse
import json
import os
import sys
from pathlib import Path

# Run from a checkout without an install: put the repo root on sys.path first.
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

from align.utils.guard import apply_resource_guards


def quality_gate_size_errors(cfg, math500_n: int, gsm8k_n: int, per_set: dict | None = None) -> list[str]:
    """Why a run of these sizes may NOT write the SFT quality-gate receipt (empty = fine).

    The receipt unlocks DPO. `assert_meets_golden` accepts any point estimate inside the target's own
    Wilson CI, which is enormous at tiny n (1/5 passes the 0.09 anchor), so the gate certifies only a
    FULL-size evaluation: both sets, at the configured n the anchors were recorded at. `per_set` (the
    measured report) is checked too, in case a set held fewer problems than requested."""
    want = {"math500": cfg.eval.math500_n, "gsm8k": cfg.eval.gsm8k_n}
    asked = {"math500": math500_n, "gsm8k": gsm8k_n}
    errs = [f"{k}: requested n={asked[k]}, the quality gate needs the full configured n={want[k]}"
            for k in want if asked[k] != want[k]]
    for k, w in want.items():
        got = (per_set or {}).get(k, {}).get("n") if per_set is not None else w
        if got != w:
            errs.append(f"{k}: measured n={got}, the quality gate needs n={w}")
    return errs


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--checkpoint", required=True, help="checkpoint dir (HF-layout safetensors dir)")
    parser.add_argument("--config", default="configs/sft.toml")
    parser.add_argument("--name", required=True, help="report model_name and output stem")
    parser.add_argument("--out", default=None, help="default: outputs/eval/<name>.json")
    parser.add_argument("--math500-n", type=int, default=None, help="0 skips; default: config")
    parser.add_argument("--gsm8k-n", type=int, default=None, help="0 skips; default: config")
    parser.add_argument(
        "--quality-gate", action="store_true", help="SFT-stage gate + ledger records"
    )
    args = parser.parse_args()

    # Guards BEFORE torch import (env vars are load-bearing); heavy=True takes
    # the exclusive job lock -- a 7B generation eval over 1,319 problems is a
    # multi-hour GPU job and must serialise behind training like any other.
    apply_resource_guards(cpu_threads=8, heavy=True, label=f"eval-{args.name}")


    from align.config import load_config  # noqa: E402
    from align.eval.golden import golden_generate_fn, load_checkpoint_model  # noqa: E402
    from align.eval.harness import EvalConfig, load_eval_set, run_eval  # noqa: E402
    from align.eval.verifier import OutcomeVerifier  # noqa: E402
    from align.tokenizer import load_tokenizer  # noqa: E402
    from align.train.stage_smoke import _manifest_hashes, _mde  # noqa: E402

    cfg = load_config(args.config)
    checkpoint = args.checkpoint
    assert os.path.isdir(checkpoint), checkpoint
    out_path = args.out or os.path.join("outputs", "eval", f"{args.name}.json")

    math500_n = cfg.eval.math500_n if args.math500_n is None else args.math500_n
    gsm8k_n = cfg.eval.gsm8k_n if args.gsm8k_n is None else args.gsm8k_n
    if args.quality_gate:
        errs = quality_gate_size_errors(cfg, math500_n, gsm8k_n)
        if errs:
            print("refusing --quality-gate on a subset eval (the receipt unlocks DPO):\n  " + "\n  ".join(errs),
                  file=sys.stderr)
            return 2
    problems = []
    # 0 in this CLI means SKIP the set (checked here); upstream load_eval_set
    # uses limit=0 for "no limit", which this branch structure never passes
    # ambiguously -- a negative value falls through to upstream validation.
    if math500_n != 0:
        problems += load_eval_set(cfg.paths.math500_path, limit=math500_n)
    if gsm8k_n != 0:
        problems += load_eval_set(cfg.paths.gsm8k_path, limit=gsm8k_n)
    if not problems:
        print("nothing to evaluate (both sets skipped)", file=sys.stderr)
        return 2

    # After every refusal above: a 7B load is the expensive step.
    tokenizer = load_tokenizer(checkpoint)
    model = load_checkpoint_model(checkpoint, device="cuda")

    verifier = OutcomeVerifier()
    gen_fn = golden_generate_fn(model, tokenizer, "cuda")
    print(f"evaluating {args.name}: {len(problems)} problems from {checkpoint}", flush=True)
    rep = run_eval(problems, gen_fn, verifier, EvalConfig(), model_name=args.name)
    per_set = rep.per_set
    for name, scores in sorted(per_set.items()):
        print(f"  {name}: n={scores.get('n')} pass@1={scores.get('pass@1')}", flush=True)

    report = {
        "per_set": per_set,
        "config": {
            "mde": _mde(cfg),
            "n_problems": sum(r.get("n", 0) for r in per_set.values()),
            "eval_set_hash": _manifest_hashes(cfg),
            "generator": "align Engine greedy (golden_generate_fn)",
            "checkpoint": checkpoint,
        },
        "pooled": rep.pooled,
        "per_problem": rep.per_problem,
    }
    os.makedirs(os.path.dirname(out_path), exist_ok=True)
    tmp = out_path + ".partial"
    with open(tmp, "w") as f:
        json.dump(report, f, indent=1, default=str)
    os.replace(tmp, out_path)
    print(f"report written: {out_path}", flush=True)

    if not args.quality_gate:
        return 0

    from align.train.stage_gates import (  # noqa: E402
        GateLedger,
        eval_source_hash,
        quality_gate,
        write_quality_gate_receipt,
    )

    errs = quality_gate_size_errors(cfg, math500_n, gsm8k_n, per_set=per_set)
    if errs:  # report is already written above; only the receipt is withheld
        print("refusing to write the quality-gate receipt:\n  " + "\n  ".join(errs), file=sys.stderr)
        return 2

    run_dir = os.path.join(cfg.paths.checkpoint_dir, "sft")
    ledger = GateLedger(run_dir, eval_source_hash=eval_source_hash())
    ledger.record("eval", passed=True, detail=json.dumps(per_set, default=str))
    decision = quality_gate(report, cfg)
    ledger.record("quality_gate", passed=decision.passed, detail=decision.halt_message[:400])
    write_quality_gate_receipt(cfg, decision, checkpoint=checkpoint)
    print(f"QUALITY GATE: passed={decision.passed}", flush=True)
    if not decision.passed:
        print(decision.halt_message, file=sys.stderr, flush=True)
    return 0 if decision.passed else 1


if __name__ == "__main__":
    # See module docstring: spawn workers re-import __main__; nothing above
    # this guard may build a model or touch the GPU.
    sys.exit(main())