"""Check that the Python package and web/router.js give the same plan.

    python tests/parity.py            # needs node on PATH

Runs `node tests/run_js.js` and router.plan on the same files and settings and compares every
number in the report: split sizes, model order, every policy's quality / lower bound / cost /
cost bound / latency / traffic shares (|difference| <= 1e-9), the cascade thresholds and router
tau, the oracle curve, and the whole sensitivity sweep. Exits non-zero on any mismatch.
"""
import json
import os
import subprocess
import sys
import tempfile

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, ROOT)
from router.data import load_csv, load_text  # noqa: E402
from router.plan import plan  # noqa: E402

TOL = 1e-9


def run_js(path, **kw):
    cmd = ["node", os.path.join(ROOT, "tests", "run_js.js"), path]
    for k, v in kw.items():
        cmd += ["--" + k, str(v)]
    res = subprocess.run(cmd, check=True, capture_output=True, cwd=ROOT)
    return json.loads(res.stdout.decode("utf-8"))


def close(a, b):
    if a is None or b is None:
        return a is None and b is None
    if isinstance(a, bool) or isinstance(b, bool):
        return a == b
    if isinstance(a, (int, float)) and isinstance(b, (int, float)):
        return abs(a - b) <= TOL * max(1.0, abs(a), abs(b))
    return a == b


def compare_eval(label, a, b, errors):
    for k in ("quality", "qualityLB", "cost", "costUB", "latency", "certified", "rowsNeeded", "tau"):
        if (k in a) != (k in b):
            errors.append(f"{label}.{k}: present in only one side")
        elif k in a and not close(a[k], b[k]):
            errors.append(f"{label}.{k}: js {a[k]!r} vs py {b[k]!r}")
    for k in ("shares", "calls", "thresholds"):
        if k in a or k in b:
            if len(a.get(k, [])) != len(b.get(k, [])) or not all(close(x, y) for x, y in zip(a[k], b[k])):
                errors.append(f"{label}.{k}: js {a.get(k)} vs py {b.get(k)}")


def compare(js, py):
    errors = []
    for k in ("n", "nCal", "nHold", "models", "signal", "binary", "enumerated", "gridSize", "aim", "requiredRate"):
        if not close(js.get(k), py.get(k)) and js.get(k) != py.get(k):
            errors.append(f"{k}: js {js.get(k)!r} vs py {py.get(k)!r}")
    for a, b in zip(js["singles"], py["singles"]):
        compare_eval("single " + a["name"], a["hold"], b["hold"], errors)
        compare_eval("single(cal) " + a["name"], a["cal"], b["cal"], errors)
    compare_eval("oracle", js["oracle"]["hold"], py["oracle"]["hold"], errors)
    ca, cb = js["oracle"]["hold"]["curve"], py["oracle"]["hold"]["curve"]
    if len(ca) != len(cb) or not all(close(p["cost"], q["cost"]) and close(p["quality"], q["quality"]) for p, q in zip(ca, cb)):
        errors.append(f"oracle curve differs ({len(ca)} vs {len(cb)} points)")
    for k in ("cascade", "router"):
        if (js.get(k) is None) != (py.get(k) is None):
            errors.append(f"{k}: present in only one side")
        elif js.get(k):
            compare_eval(k, js[k]["hold"], py[k]["hold"], errors)
            compare_eval(k + "(cal)", js[k]["cal"], py[k]["cal"], errors)
            if js[k].get("feasible") != py[k].get("feasible"):
                errors.append(f"{k}.feasible: js {js[k].get('feasible')} vs py {py[k].get('feasible')}")
    sa, sb = js["sensitivity"], py["sensitivity"]
    if [r["target"] for r in sa] != [r["target"] for r in sb]:
        errors.append(f"sensitivity targets: js {[r['target'] for r in sa]} vs py {[r['target'] for r in sb]}")
    else:
        for ra, rb in zip(sa, sb):
            for k in ("cascade", "router"):
                if (k in ra) != (k in rb):
                    errors.append(f"sensitivity {ra['target']}.{k}: present in only one side")
                elif k in ra:
                    compare_eval(f"sensitivity {ra['target']}.{k}", ra[k], rb[k], errors)
            if not close(ra.get("oracleCost"), rb.get("oracleCost")):
                errors.append(f"sensitivity {ra['target']}.oracleCost: js {ra.get('oracleCost')} vs py {rb.get('oracleCost')}")
            if (ra.get("bestSingle") or {}).get("name") != (rb.get("bestSingle") or {}).get("name"):
                errors.append(f"sensitivity {ra['target']}.bestSingle: js {ra.get('bestSingle')} vs py {rb.get('bestSingle')}")
    if js.get("breakEven") != py.get("breakEven") and not (
            js.get("breakEven") and py.get("breakEven") and all(close(js["breakEven"][k], py["breakEven"][k]) for k in js["breakEven"])):
        errors.append(f"breakEven: js {js.get('breakEven')} vs py {py.get('breakEven')}")
    return errors


def to_long(path, out):
    """Rewrite the wide demo as the long layout, with a score column instead of 0/1 and no latency."""
    import csv
    with open(path, encoding="utf-8", newline="") as f:
        rows = list(csv.DictReader(f))
    models = sorted({k[:-8] for k in rows[0] if k.endswith("_correct")})
    with open(out, "w", encoding="utf-8", newline="") as f:
        w = csv.writer(f, delimiter=";")
        w.writerow(["id", "model", "score", "cost_usd", "p_correct", "question"])
        for r in rows:
            for mo in models:
                score = r[mo + "_correct"]
                if r["prompt_id"].endswith("7"):
                    score = "0.5"                  # a few partial credits make it a score, not 0/1
                w.writerow([r["prompt_id"], mo, score, r[mo + "_cost"], r[mo + "_confidence"], r["prompt"]])


def no_conf(path, out):
    import csv
    with open(path, encoding="utf-8", newline="") as f:
        rd = csv.DictReader(f)
        cols = [c for c in rd.fieldnames if not c.endswith("_confidence")]
        rows = list(rd)
    with open(out, "w", encoding="utf-8", newline="") as f:
        w = csv.DictWriter(f, fieldnames=cols, extrasaction="ignore")
        w.writeheader()
        w.writerows(rows[:600])


def main():
    try:
        subprocess.run(["node", "--version"], check=True, capture_output=True)
    except (OSError, subprocess.CalledProcessError):
        print("node is not available; can't run the parity check")
        return 2
    demo = os.path.join(ROOT, "web", "demo.csv")
    failed = 0
    with tempfile.TemporaryDirectory() as tmp:
        longf, nocf = os.path.join(tmp, "long.csv"), os.path.join(tmp, "noconf.csv")
        to_long(demo, longf)
        no_conf(demo, nocf)
        cases = [
            ("demo, target 0.9", demo, dict(target=0.9, delta=0.05)),
            ("demo, target 0.85, delta 0.01, cal 0.6, seed 3", demo, dict(target=0.85, delta=0.01, cal=0.6, seed=3)),
            ("demo, learned signal", demo, dict(target=0.88, signal="learned")),
            ("long layout, scores in [0,1], ';' delimiter", longf, dict(target=0.8)),
            ("wide without confidence (learned router only), 600 rows", nocf, dict(target=0.85, safety=0)),
        ]
        for label, path, kw in cases:
            js = run_js(path, **kw)
            log = load_csv(path)
            py = json.loads(json.dumps(plan(log, target=kw.get("target", 0.9), delta=kw.get("delta", 0.05),
                                            cal_frac=kw.get("cal", 0.5), seed=kw.get("seed", 1),
                                            signal=kw.get("signal", "auto"), safety=kw.get("safety", 1))))
            errors = compare(js, py)
            print(("ok   " if not errors else "FAIL ") + label + f"   (js {js['ms']} ms)")
            for e in errors[:20]:
                print("     " + e)
            failed += bool(errors)
    print(f"{len(cases) - failed}/{len(cases)} cases identical")
    return 1 if failed else 0


if __name__ == "__main__":
    sys.exit(main())