"""Model proof: one call to the Sentinel model completes a multi-step job, and we check the result ourselves.

Each job is sent as a SINGLE OpenAI-style request to POST /v1/chat/completions (model `sentinel`, the default
open-weight base). Nothing else is sent: no tools are defined by the caller, no follow-up messages. Afterwards the
harness verifies the work independently (hidden checks the model never sees) and reads the session transcript for the
tool sequence, steps, tokens and any self-heal notices.

Run on the machine that serves the model (the job folders must be visible to it):
    .venv/bin/python evals/model_proof.py [--runs 3] [--url http://127.0.0.1:8152] [--only bugfix,invoices,logs]
Writes evals/results/model-proof-<id>.json (one row per call + a summary) and prints the summary.
"""
from __future__ import annotations

import argparse
import hashlib
import json
import os
import random
import re
import subprocess
import sys
import tempfile
import time
import uuid
from pathlib import Path

import httpx

LAB = Path(__file__).resolve().parents[1]
TOKEN = next((l.split("=", 1)[1].strip() for l in open(LAB / ".env") if l.startswith("SENTINEL_TOKEN=")), "")


# ───────────── jobs: fixtures, the one prompt, and the hidden check ─────────────
def job_bugfix(d: Path, write: bool = True) -> dict:
    TESTS = ("import unittest\nfrom stats import mean, median, spread\n\n\n"
             "class T(unittest.TestCase):\n"
             "    def test_mean(self):\n        self.assertEqual(mean([1, 2, 3, 4]), 2.5)\n\n"
             "    def test_median_odd(self):\n        self.assertEqual(median([5, 1, 3]), 3)\n\n"
             "    def test_median_even(self):\n        self.assertEqual(median([4, 1, 3, 2]), 2.5)\n\n"
             "    def test_spread(self):\n        self.assertEqual(spread([3, 9, 4]), 6)\n\n\n"
             "if __name__ == '__main__':\n    unittest.main()\n")
    write and (d / "stats.py").write_text(
        "def mean(xs):\n    return sum(xs) / len(xs)\n\n\n"
        "def median(xs):\n    s = sorted(xs)\n    n = len(s)\n    return s[n // 2]\n\n\n"
        "def spread(xs):\n    return max(xs) - min(xs)\n")
    write and (d / "test_stats.py").write_text(TESTS)
    tests_hash = hashlib.sha256(TESTS.encode()).hexdigest()
    prompt = (f"The unit tests in {d} fail. Find the bug, fix the code (do not edit the tests), run the tests with "
              f"`python3 -m unittest` in that folder until they pass, then reply with one sentence saying what was wrong.")

    def check(answer: str) -> tuple[bool, list[str]]:
        why = []
        if hashlib.sha256((d / "test_stats.py").read_bytes()).hexdigest() != tests_hash:
            why.append("the tests were edited")
        r = subprocess.run([sys.executable, "-m", "unittest", "-q"], cwd=d, capture_output=True, text=True, timeout=60)
        if r.returncode != 0:
            why.append("tests still fail")
        hidden = subprocess.run([sys.executable, "-c", "from stats import median as m; assert m([10,2,8,4])==6; assert m([7])==7; assert m([3,1,2])==2; assert m([1,2,3,4,5,6])==3.5"],
                                cwd=d, capture_output=True, text=True, timeout=30)
        if hidden.returncode != 0:
            why.append("hidden median cases fail")
        if not re.search(r"median|even", answer, re.I):
            why.append("the reply does not say what was wrong")
        return not why, why
    return {"prompt": prompt, "check": check}


def job_invoices(d: Path, write: bool = True) -> dict:
    rows = [("Acme", "1,200.50"), ("Borealis", "310.00"), ("Acme", "89.99"), ("Cobalt", "2,045.10"), ("Borealis", "42.50"),
            ("Dune", "999.99"), ("Cobalt", "15.00"), ("Acme", "430.01"), ("Dune", "1,000.00"), ("Borealis", "7.25")]
    if write:
        (d / "invoices.csv").write_text("customer,amount\n" + "".join(f'{c},"{a}"\n' for c, a in rows))
    expected: dict[str, float] = {}
    for c, a in rows:
        expected[c] = round(expected.get(c, 0) + float(a.replace(",", "")), 2)
    grand = f"{sum(expected.values()):,.2f}"
    prompt = (f"In {d} there is invoices.csv. Write totals.py with a function totals(path) that returns a dict of customer -> "
              f"total amount (amounts are quoted and can contain thousands separators). Write test_totals.py with unittest "
              f"that checks it against that file, run `python3 -m unittest` in that folder until it passes, then reply with "
              f"the grand total of all invoices to two decimals.")

    def check(answer: str) -> tuple[bool, list[str]]:
        why = []
        if not (d / "totals.py").exists():
            return False, ["totals.py missing"]
        probe = subprocess.run([sys.executable, "-c", f"import json; from totals import totals; print(json.dumps({{str(k): float(v) for k, v in totals({str(d / 'invoices.csv')!r}).items()}}))"],
                               cwd=d, capture_output=True, text=True, timeout=30)
        try:
            got = {k: round(float(v), 2) for k, v in json.loads(probe.stdout.strip().splitlines()[-1]).items()}
            if got != expected:
                why.append(f"totals() wrong: {got}")
        except Exception:  # noqa: BLE001
            why.append("totals() did not run: " + (probe.stderr.strip().splitlines() or ["?"])[-1][:120])
        if not (d / "test_totals.py").exists():
            why.append("test_totals.py missing")
        elif subprocess.run([sys.executable, "-m", "unittest", "-q"], cwd=d, capture_output=True, timeout=60).returncode != 0:
            why.append("its own tests fail")
        if grand not in answer and grand.replace(",", "") not in answer:
            why.append(f"grand total {grand} not in the reply")
        return not why, why
    return {"prompt": prompt, "check": check, "expected": grand}


def job_logs(d: Path, write: bool = True) -> dict:
    rnd = random.Random(1014)
    codes = {"E101": 41, "E204": 27, "E330": 19, "E412": 8, "E505": 3}
    lines = [f"2026-10-14T0{rnd.randint(0, 9)}:{rnd.randint(10, 59)}:{rnd.randint(10, 59)} INFO request ok id={i}" for i in range(220)]
    for code, n in codes.items():
        lines += [f"2026-10-14T0{rnd.randint(0, 9)}:{rnd.randint(10, 59)}:{rnd.randint(10, 59)} ERROR {code} worker={rnd.randint(1, 9)}" for _ in range(n)]
    rnd.shuffle(lines)
    write and (d / "app.log").write_text("\n".join(lines) + "\n")
    prompt = (f"Analyse {d}/app.log. Write {d}/report.md with a markdown table of every ERROR code and how many times it "
              f"appears, most frequent first. Then reply with the most frequent code and its count.")

    def check(answer: str) -> tuple[bool, list[str]]:
        why = []
        rep = d / "report.md"
        if not rep.exists():
            return False, ["report.md missing"]
        text = rep.read_text()
        rows = [(m.group(1), int(m.group(2))) for m in re.finditer(r"\|\s*(E\d{3})\s*\|\s*(\d+)\s*\|", text)]
        if sorted(rows) != sorted(codes.items()):
            why.append(f"table counts wrong: {rows}")
        elif [c for c, _ in rows] != [c for c, _ in sorted(codes.items(), key=lambda x: -x[1])]:
            why.append("table not sorted most frequent first")
        if "E101" not in answer or "41" not in answer:
            why.append("reply does not name E101 with 41")
        return not why, why
    return {"prompt": prompt, "check": check}


JOBS = {"bugfix": job_bugfix, "invoices": job_invoices, "logs": job_logs}


# ───────────── one call, then read the transcript ─────────────
def run_one(c: httpx.Client, name: str, model: str, root: Path) -> dict:
    d = root / f"{name}-{uuid.uuid4().hex[:6]}"
    d.mkdir(parents=True)
    job = JOBS[name](d)
    t0 = time.time()
    r = c.post("/v1/chat/completions", json={"model": model, "messages": [{"role": "user", "content": job["prompt"]}]})
    secs = round(time.time() - t0, 1)
    row = {"job": name, "folder": str(d), "prompt": job["prompt"], "seconds": secs, "http": r.status_code}
    if r.status_code != 200:
        row.update(ok=False, why=[f"HTTP {r.status_code}: {r.text[:200]}"]); return row
    body = r.json()
    answer = body["choices"][0]["message"]["content"]
    ok, why = job["check"](answer)
    meta = body.get("sentinel", {})
    s = c.get(f"/api/sessions/{meta.get('session_id')}").json()
    tools = [{"name": e["name"], "seconds": e.get("seconds")} for e in s.get("events", []) if e.get("type") == "tool_result" and e.get("agent") is not None]
    calls = [e["name"] for e in s.get("events", []) if e.get("type") == "tool_call"]
    heals = [e["text"] for e in s.get("events", []) if e.get("type") == "notice"]
    host = meta.get("base_host") or next((m.get("host", "") for m in json.loads((LAB / "models.json").read_text()) if m.get("name") == meta.get("base_model")), "")
    row.update(ok=ok, why=why, answer=answer[:600], base_model=meta.get("base_model"), base_host=host, steps=meta.get("steps"), stopped=meta.get("stopped"),
               tool_calls=calls, tool_seconds=round(sum(t["seconds"] or 0 for t in tools), 1), self_heals=heals,
               prompt_tokens=body["usage"]["prompt_tokens"], completion_tokens=body["usage"]["completion_tokens"], session_id=meta.get("session_id"))
    return row


def summarize(rows: list[dict], names: list[str], started: float | None = None, date: str | None = None) -> dict:
    passed = sum(r["ok"] for r in rows)
    return {"calls": len(rows), "passed": passed, "pass_rate": round(passed / max(len(rows), 1), 3),
            "median_seconds": sorted(r["seconds"] for r in rows)[len(rows) // 2],
            "median_tool_calls": sorted(len(r.get("tool_calls") or []) for r in rows)[len(rows) // 2],
            "base_models": sorted({r.get("base_model") for r in rows if r.get("base_model")}),
            "base_hosts": sorted({r.get("base_host") for r in rows if r.get("base_host")}),
            "by_job": {n: f"{sum(r['ok'] for r in rows if r['job'] == n)}/{sum(r['job'] == n for r in rows)}" for n in names},
            "rechecked": sum(bool(r.get("rechecked")) for r in rows),
            **({"wall_minutes": round((time.time() - started) / 60, 1)} if started else {}), "date": date or time.strftime("%Y-%m-%d")}


def recheck(path: Path, reason: str) -> int:
    """Re-evaluate every row of a results file against the folder the model left behind (fixtures are NOT rewritten).
    Used when a check itself was wrong; each changed row keeps its original verdict and the reason, so nothing is hidden."""
    res = json.loads(path.read_text())
    for r in res["rows"]:
        job = JOBS[r["job"]](Path(r["folder"]), write=False)
        ok, why = job["check"](r.get("answer") or "")
        if ok != r["ok"] or why != r["why"]:
            r.setdefault("original", {"ok": r["ok"], "why": r["why"]})
            r.update(ok=ok, why=why, rechecked=reason)
            print(f"[{r['job']} #{r['run']}] {r['original']['ok']} -> {ok} {why}")
    names = list(dict.fromkeys(r["job"] for r in res["rows"]))
    res["summary"] = {**res.get("summary", {}), **summarize(res["rows"], names, date=(res.get("summary") or {}).get("date"))}
    path.write_text(json.dumps(res, indent=1))
    print(json.dumps(res["summary"], indent=1))
    return 0


def main() -> int:
    if "--recheck" in sys.argv:
        i = sys.argv.index("--recheck")
        return recheck(Path(sys.argv[i + 1]), sys.argv[i + 2] if len(sys.argv) > i + 2 else "check corrected after the run")
    ap = argparse.ArgumentParser()
    ap.add_argument("--url", default=os.environ.get("SENTINEL_URL", "http://127.0.0.1:" + os.environ.get("SENTINEL_PORT", "8152")))
    ap.add_argument("--model", default="sentinel")
    ap.add_argument("--runs", type=int, default=3)
    ap.add_argument("--only", default="")
    a = ap.parse_args()
    names = [n for n in (a.only.split(",") if a.only else JOBS)]
    c = httpx.Client(base_url=a.url, timeout=3600, headers={"Authorization": "Bearer " + TOKEN})
    root = Path(tempfile.mkdtemp(prefix="model-proof-"))
    run_id = time.strftime("%Y%m%d-%H%M") + "-" + uuid.uuid4().hex[:4]
    out = LAB / "evals" / "results" / f"model-proof-{run_id}.json"
    out.parent.mkdir(parents=True, exist_ok=True)
    rows: list[dict] = []
    started = time.time()
    for i in range(a.runs):
        for name in names:
            row = run_one(c, name, a.model, root)
            row["run"] = i + 1
            rows.append(row)
            print(f"[{row['job']} #{i + 1}] {'PASS' if row['ok'] else 'FAIL'} {row['seconds']}s steps={row.get('steps')} tools={row.get('tool_calls')} {row['why']}", flush=True)
            out.write_text(json.dumps({"run_id": run_id, "url": a.url, "model": a.model, "rows": rows}, indent=1))
    summary = summarize(rows, names, started)
    out.write_text(json.dumps({"run_id": run_id, "url": a.url, "model": a.model, "summary": summary, "rows": rows}, indent=1))
    print(json.dumps(summary, indent=1))
    print("wrote", out)
    return 0


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