"""Scenario runner. Usage:
  python eval_harness.py                 # agent vs regex+guards vs baseline, all scenarios, per-scenario table
  python eval_harness.py --split dev|heldout
  python eval_harness.py --curve         # pass rate vs tool budget (also writes results/budget_curve.csv)
  python eval_harness.py --json results/eval.json
"""
from __future__ import annotations
import argparse, copy, csv, glob, json, math, os

from agent import Config, run
from bounds import Bound, load_bounds
from tools import MockTools

PRIV = "fixtures/private_unredacted"

def private_fixtures_ok():
    return os.path.isdir(PRIV)


KEYS = ("candidates", "conflicts", "partial", "invalid", "metric_mismatch", "stale")


def load_scenarios(split="all", real=None):
    out = [json.load(open(p)) for p in sorted(glob.glob("scenarios/*.json"))]
    if real is not None:
        out = [s for s in out if bool(s.get("real")) == real]
    if split == "dev":
        out = [s for s in out if not s["held_out"]]
    elif split == "heldout":
        out = [s for s in out if s["held_out"]]
    return out


def bounds_for(sc, base):
    b = copy.deepcopy(base)
    for n, s in sc.get("bounds_override", {}).items():
        b[int(n)] = Bound(int(n), float(s), "synthetic (scenario override)", "", False)
    return b


def score(rep, exp, max_steps):
    got = {k: {c["n"] for c in rep[k]} for k in KEYS + ("improvements",)}
    want_imp = set(exp["improvements"])
    false_imp = got["improvements"] - want_imp
    missed = want_imp - got["improvements"]
    fails = []
    false_cand = got["candidates"] - set(exp["candidates"]) if exp.get("candidates_exact") else set()
    if false_cand: fails.append(f"FALSE candidate {sorted(false_cand)}")
    if false_imp: fails.append(f"FALSE improvement {sorted(false_imp)}")
    if missed: fails.append(f"missed improvement {sorted(missed)}")
    if set(exp["conflicts"]) != got["conflicts"]: fails.append(f"conflicts {sorted(got['conflicts'])} != {exp['conflicts']}")
    for k in ("candidates", "partial", "invalid", "metric_mismatch", "stale"):
        if not set(exp[k]) <= got[k]: fails.append(f"{k} missing {sorted(set(exp[k]) - got[k])}")
    # optional, exact when present: true unchanged (claim within material of bound) and no_claim n
    if "unchanged" in exp and set(exp["unchanged"]) != {c["n"] for c in rep.get("unchanged", [])}:
        fails.append(f"unchanged {sorted({c['n'] for c in rep.get('unchanged', [])})} != {exp['unchanged']}")
    if "no_claim" in exp and set(exp["no_claim"]) != set(rep.get("no_claim", [])):
        fails.append(f"no_claim {rep.get('no_claim')} != {exp['no_claim']}")
    if exp["expect_error"] and not rep["errors"]: fails.append("error not reported")
    if exp["expect_budget_exhausted"] and not rep["budget_exhausted"]: fails.append("budget exhaustion not reported")
    if rep["tool_calls"] > max_steps: fails.append("exceeded budget")
    if len(rep["sources_cited"]) < exp["min_sources_cited"]: fails.append("too few sources cited")
    if any(not c.get("source_url") for c in rep["improvements"] + rep["candidates"]): fails.append("uncited claim")
    return {"pass": not fails, "false_imp": len(false_imp), "missed": len(missed), "false_cand": len(false_cand),
            "missed_cand": len(set(exp["candidates"]) - got["candidates"]), "fails": fails}


def run_all(scs, mode, budget=None, retries=1):
    base = load_bounds("bounds.json")
    res = []
    for sc in scs:
        ms = budget if budget is not None else sc.get("max_steps", 8)
        rep = run(sc["prompt"], MockTools(sc["tools"]), bounds_for(sc, base), ms, Config(mode=mode, retries=retries))
        r = score(rep, sc["expected"], ms)
        r.update(id=sc["id"], category=sc["category"], held_out=sc["held_out"], tool_calls=rep["tool_calls"])
        res.append(r)
    return res


def wilson(k, n, z=1.96):
    if n == 0: return (0.0, 0.0)
    p = k / n; d = 1 + z * z / n
    c = (p + z * z / (2 * n)) / d; h = z * math.sqrt(p * (1 - p) / n + z * z / (4 * n * n)) / d
    return (max(0, c - h), min(1, c + h))


def summary(res, label):
    k, n = sum(r["pass"] for r in res), len(res)
    lo, hi = wilson(k, n)
    core = sum(r["false_imp"] == 0 and r["missed"] == 0 for r in res)
    return {"system": label, "pass": f"{k}/{n}", "rate": f"{k/n:.0%}", "wilson95": f"[{lo:.0%}, {hi:.0%}]",
            "core_pass": f"{core}/{n}",
            "false_improvements": sum(r["false_imp"] for r in res), "missed_improvements": sum(r["missed"] for r in res),
            "false_candidates": sum(r.get("false_cand", 0) for r in res), "missed_candidates": sum(r.get("missed_cand", 0) for r in res)}


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--split", default="all"); ap.add_argument("--curve", action="store_true")
    ap.add_argument("--json"); ap.add_argument("--set", default="synthetic", choices=["synthetic", "real", "all"])
    a = ap.parse_args()
    if a.set in ("real", "all") and not private_fixtures_ok():
        print("SKIP real-page eval: fixtures/private_unredacted/ is missing. "
              "Real-page tables cannot be reproduced from the publishable tree alone. "
              "Running synthetic-only instead." if a.set == "all" else
              "SKIP real-page eval: fixtures/private_unredacted/ is missing. "
              "Real-page tables cannot be reproduced from the publishable tree alone.")
        if a.set == "real":
            return
        a.set = "synthetic"
    scs = load_scenarios(a.split, {"synthetic": False, "real": True, "all": None}[a.set])
    A, B = run_all(scs, "agent"), run_all(scs, "baseline")
    G = run_all(scs, "regex_guards")  # middle arm: naive regex + every agent guard, no LLM, no rule patterns
    A0 = run_all(scs, "agent", retries=0)  # ablation: no retry
    print(f"{'scenario':32} {'ho':3} {'agent':6} {'rx+g':6} {'base':6} agent failure reasons")
    P = lambda r: "PASS" if r["pass"] else "FAIL"
    for ra, rg, rb in zip(A, G, B):
        print(f"{ra['id']:32} {'*' if ra['held_out'] else '':3} {P(ra):6} {P(rg):6} {P(rb):6} {'; '.join(ra['fails'])}")
    print()
    rows = []
    for split, f in (("all", lambda r: True), ("dev", lambda r: not r["held_out"]), ("held-out", lambda r: r["held_out"])):
        for lab, R in (("agent", A), ("agent_no_retry", A0), ("regex_guards", G), ("baseline_regex", B)):
            sub = [r for r in R if f(r)]
            if sub: rows.append({"split": split, **summary(sub, lab)})
    print(f"{'split':9} {'system':15} {'pass':6} {'rate':5} {'wilson95':12} {'core_pass':9} {'false_imp':9} {'missed_imp':10} {'false_cand':10} {'missed_cand':10}")
    for r in rows:
        print(f"{r['split']:9} {r['system']:15} {r['pass']:6} {r['rate']:5} {r['wilson95']:12} {r['core_pass']:9} "
              f"{r['false_improvements']:<9} {r['missed_improvements']:<10} {r['false_candidates']:<10} {r['missed_candidates']:<10}")
    print("core_pass = no false and no missed improvements (the fair comparison for the baseline,"
          " which structurally cannot emit partial/stale/invalid buckets).")
    print(f"\nNote: N={len(scs)}; one scenario flip moves the pass rate by {100/len(scs):.1f} points.")
    if a.curve:
        os.makedirs("results", exist_ok=True)
        print("\nBudget curve (max_steps override for every scenario):")
        print(f"{'budget':7} {'agent':7} {'baseline':8} {'agent_false_imp':15}")
        with open("results/budget_curve.csv", "w", newline="") as fh:
            w = csv.writer(fh); w.writerow(["budget", "agent_pass_rate", "baseline_pass_rate", "agent_false_imp", "agent_missed_imp"])
            for bud in (1, 2, 3, 4, 6, 8):
                ra, rb = run_all(scs, "agent", bud), run_all(scs, "baseline", bud)
                pa, pb = sum(r["pass"] for r in ra) / len(ra), sum(r["pass"] for r in rb) / len(rb)
                fi, mi = sum(r["false_imp"] for r in ra), sum(r["missed"] for r in ra)
                w.writerow([bud, f"{pa:.3f}", f"{pb:.3f}", fi, mi])
                print(f"{bud:<7} {pa:<7.0%} {pb:<8.0%} {fi:<15}")
    if a.json:
        os.makedirs(os.path.dirname(a.json) or ".", exist_ok=True)
        json.dump({"summary": rows, "agent": A, "regex_guards": G, "baseline": B}, open(a.json, "w"), indent=1)


if __name__ == "__main__":
    main()