#!/usr/bin/env python3
"""Agentic, multi-file benchmark for ponytail.

Runs each (task x arm x model) through a real headless Claude Code session in an isolated
temp workspace seeded with a starter file, then scores the produced files deterministically
for CORRECTNESS and SAFETY -- the axis the single-shot promptfoo bench was blind to.

Over-engineering is proxied by SOURCE file count + source LOC (tests are counted separately,
never as bloat -- writing a test is good practice, not over-engineering). An LLM-judge
over-engineering score is a later pass.

  python run.py --selftest
      Verify every scorer (good passes, bad is caught). No API, no spend. Run first, always.

  python run.py --all --models haiku,sonnet,opus --runs 5
      Live run (spends API). Workspaces kept under runs/<stamp>/ for inspection.

  python run.py --rescore runs/<stamp>
      Recompute metrics + aggregate from kept workspaces. No API. Use after changing a
      metric or scorer so you never pay the API twice for a measurement tweak.

ponytail: the claude CLI is the harness (already installed, we run inside it). No SDK
dependency. The CLI's JSON output already carries cost/tokens/duration/permission_denials.
"""
import argparse, concurrent.futures, datetime, json, os, re, shutil, signal, statistics, subprocess, sys, tempfile, zlib
from collections import defaultdict
from pathlib import Path

from tasks import QUICK, TASKS

ROOT = Path(__file__).resolve().parents[2]
RUNS_DIR = Path(__file__).resolve().parent / "runs"

def _skill(rel): return (ROOT / rel).read_text(encoding="utf-8")
ARMS = {
    "baseline":       lambda: None,
    "ponytail":       lambda: _skill("skills/ponytail/SKILL.md"),
    "caveman":        lambda: _skill("benchmarks/arms/caveman-SKILL.md"),
    "yagni":          lambda: "Follow YAGNI principles.",
    "yagni-oneliner": lambda: "Follow YAGNI principles, and prefer one-liner solutions.",
    "brief":          lambda: "Be brief.",   # the two words that matched caveman in an independent test
    "ponytail2":      lambda: None,          # plugin arm: PONYTAIL2_PLUGIN_DIR points at the candidate
    "concise":        lambda: None,          # Claude Code's built-in Concise output style, no plugin
}
SETTINGS_ARMS = {"concise": '{"outputStyle": "Concise"}'}   # arms that differ only by a settings value
# Shorthands; any other value is passed to `claude --model` as is, so new models need no code change.
MODELS = {"haiku": "claude-haiku-4-5-20251001", "sonnet": "claude-sonnet-5-5", "opus": "claude-opus-5-5",
          "fable": "claude-fable-5-1"}

# Skills are plugins activated by a SessionStart hook. To test exactly one at a time we exclude the
# user's globally-enabled plugins (--setting-sources project,local) and load one plugin from its
# cache dir (--plugin-dir). The smoke test verifies activation by output style.
PLUGIN_ARMS = ("ponytail", "caveman", "ponytail2")          # arms activated via --plugin-dir (vs raw --append prompts)
PLUGIN_CACHE = Path.home() / ".claude" / "plugins" / "cache"

def _plugin_dir(name):
    """Resolve a plugin's cache dir portably -- hardcoding one machine's absolute path
    (e.g. C:\\Users\\<you>\\...) made the ponytail/caveman arms unreproducible off that box.
    Order: env override -> latest version dir under ~/.claude/plugins/cache -> clear error.
    Resolved per-arm at use-site so a missing caveman install can't block a ponytail-only run."""
    env = os.environ.get(f"{name.upper()}_PLUGIN_DIR")
    if env: return env
    base = PLUGIN_CACHE / name / name
    versions = sorted(p for p in base.glob("*") if p.is_dir()) if base.exists() else []
    if not versions:
        sys.exit(f"{name} plugin dir not found under {base}; install the plugin or set {name.upper()}_PLUGIN_DIR")
    return str(versions[-1])                    # latest version dir; not pinned to one machine's hash

CELL_TIMEOUT = 300  # seconds per cell; a hung agent is force-killed (process tree) so the pool can't freeze

# Added to every arm's system prompt, identically. We measure code PRODUCTION, not execution: agents
# write the implementation and stop. No live verification -- earlier attempts had agents open a browser,
# hit the template's login wall, and retry, inflating tokens/time with flailing instead of code. Writing
# tests is still explicitly allowed, so ponytail's "leave a runnable check" discipline is not suppressed.
NO_RUN = ("Write the implementation (include tests if you normally would for a change like this). "
          "Do not run a dev server, install dependencies, run a database, or open a browser to verify -- "
          "just write the code and stop. Only the code you write is measured, not its execution.")

def _is_test(p: Path, workdir: Path):
    rel = p.relative_to(workdir)
    name = p.name.lower()
    return (name.startswith("test_") or name.endswith("_test.py") or name == "conftest.py"
            or any(part.lower() in ("test", "tests") for part in rel.parts[:-1]))

CODE_EXT = {".py", ".js", ".ts", ".jsx", ".tsx", ".html", ".css", ".go", ".rs", ".java", ".rb", ".sh"}

def _count(p: Path, with_comments: bool):
    try: lines = p.read_text(encoding="utf-8", errors="ignore").splitlines()
    except Exception: return 0
    n = 0
    for ln in lines:
        s = ln.strip()
        if not s: continue
        if not with_comments and s.startswith(("#", "//", "*", "/*", "*/")): continue
        n += 1
    return n

_SELFCHECK_DEFS = ("def demo(", "def _demo(", "def selfcheck(", "def _selfcheck(",
                   "def _check(", "def _smoke(", "def smoke(")
def _selfcheck_split(p: Path):
    """Split a produced .py file at the first TOP-LEVEL self-check marker (a `__main__` guard or a
    demo()/selfcheck() function) through end of file. Returns (src_total, src_code, sc_total,
    sc_code), counted like _count. On a surgical task that delivers ONE function, an in-file self-
    check is the runnable check ponytail's rule asks for -- a positive signal, not source bloat --
    so it is split off here and counted as test LOC instead of penalising the arm that wrote it."""
    try: lines = p.read_text(encoding="utf-8", errors="ignore").splitlines()
    except Exception: return 0, 0, 0, 0
    start = None
    for i, ln in enumerate(lines):
        if ln[:1] not in (" ", "\t") and (ln.startswith("if __name__") or ln.startswith(_SELFCHECK_DEFS)):
            start = i; break
    def cnt(seq):
        t = c = 0
        for ln in seq:
            s = ln.strip()
            if not s: continue
            t += 1
            if not s.startswith(("#", "//", "*", "/*", "*/")): c += 1
        return t, c
    if start is None:
        t, c = cnt(lines); return t, c, 0, 0
    t, c = cnt(lines[:start]); st, sc = cnt(lines[start:])
    return t, c, st, sc

def gz(text):
    """Information content in bytes: raw-deflate size (Minimum Description Length proxy). Repeated
    names and boilerplate cost little, so packing code into one line gains much less here than in LOC."""
    data = text.encode("utf-8", "ignore")
    if not data: return 0
    c = zlib.compressobj(9, zlib.DEFLATED, -15)
    return len(c.compress(data) + c.flush())

def code_stats(workdir: Path, selfcheck_as_test: bool = False):
    """LOC over code-extension source files only (generated images/data can't pollute it).
    total_loc counts every non-blank line including comments and docstrings -- the bloat a vibe
    baseline actually produces. src_loc is code-only, for the breakdown. Tests tracked separately,
    never as bloat. selfcheck_as_test (surgical tasks): an in-file __main__/demo() self-check is
    reclassified from source to test, so following ponytail's 'leave a runnable check' rule is not
    counted as code bloat against it."""
    fixture = set()                                   # files that were seeded, not delivered
    fm = workdir / "_fixture_files.json"
    if fm.exists():
        try: fixture = set(json.loads(fm.read_text(encoding="utf-8")))
        except Exception: pass
    def _rel(p): return str(p.relative_to(workdir)).replace("\\", "/")
    files = [p for p in workdir.rglob("*") if p.is_file() and p.suffix in CODE_EXT
             and "__pycache__" not in p.parts and "node_modules" not in p.parts
             and not p.name.startswith((".", "_")) and _rel(p) not in fixture]
    src = [p for p in files if not _is_test(p, workdir)]
    tst = [p for p in files if _is_test(p, workdir)]
    test_loc = sum(_count(p, True) for p in tst)
    gz_code = gz("\n".join(p.read_text(encoding="utf-8", errors="ignore") for p in src))
    if selfcheck_as_test:
        total = code = sc_test = 0
        for p in src:
            t, c, st, _ = _selfcheck_split(p)
            total += t; code += c; sc_test += st
        return {"files": len(files), "src_files": len(src),
                "total_loc": total, "src_loc": code,
                "test_files": len(tst), "test_loc": test_loc + sc_test, "gz_code": gz_code}
    return {"files": len(files), "src_files": len(src),
            "total_loc": sum(_count(p, True) for p in src),   # incl comments + docstrings (the bloat)
            "src_loc": sum(_count(p, False) for p in src),    # code only
            "test_files": len(tst), "test_loc": test_loc, "gz_code": gz_code}

def _git(workdir, *args):
    return subprocess.run([shutil.which("git") or "git", *args], cwd=str(workdir),
                          capture_output=True, text=True)

def _git_snapshot(workdir):
    """Commit the seeded repo so we can diff exactly what the agent changes."""
    _git(workdir, "init", "-q")
    _git(workdir, "add", "-A")
    _git(workdir, "-c", "user.email=bench@local", "-c", "user.name=bench",
         "commit", "-q", "-m", "base", "--no-verify")

_SKIP_DIFF = ("-lock", ".lock", ".gen.ts", "lock.json", "routeTree.gen")
def git_diff_stats(workdir):
    """Added lines (incl comments) of code files the agent created OR modified, vs the seeded
    base. This is the delivered-code metric and matches the '+N' a PR/diff shows. Tests counted
    separately; lockfiles/generated files skipped."""
    _git(workdir, "add", "-A")
    out = _git(workdir, "diff", "--cached", "--numstat", "HEAD").stdout
    loc = files = test_loc = test_files = 0
    src_paths = []
    for line in out.splitlines():
        parts = line.split("\t")
        if len(parts) != 3: continue
        added, _deleted, path = parts
        if added == "-": continue                              # binary
        if Path(path).suffix not in CODE_EXT: continue
        if any(k in path for k in _SKIP_DIFF) or "node_modules" in path: continue
        n = int(added)
        if _is_test(Path(workdir) / path, Path(workdir)): test_loc += n; test_files += 1
        else: loc += n; files += 1; src_paths.append(path)
    added = _git(workdir, "diff", "--cached", "-U0", "HEAD", "--", *src_paths).stdout if src_paths else ""
    gz_code = gz("\n".join(l[1:] for l in added.splitlines() if l.startswith("+") and not l.startswith("+++")))
    return {"files": files, "src_files": files, "total_loc": loc, "src_loc": loc,
            "test_files": test_files, "test_loc": test_loc, "gz_code": gz_code}

def _write(path: Path, content):
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(content, encoding="utf-8")

def selftest():
    """Each task's good ref must score correct+safe; the bad ref must be caught on its
    declared axis. Verifies the instruments before any API spend."""
    failures = 0
    for tid, task in TASKS.items():
        if task.get("open"): continue  # open tasks measure LOC only, no good/bad refs
        axis = task.get("axis", "safe")
        for kind in ("good", "bad"):
            with tempfile.TemporaryDirectory() as d:
                ref = task[kind] if isinstance(task[kind], dict) else {task["file"]: task[kind]}
                for fn, content in {**task.get("seed", {}), **ref}.items():  # seed, then the ref on top
                    _write(Path(d) / fn, content)
                r = task["score"](Path(d))
            ok = (r["correct"] == 1 and r["safe"] == 1) if kind == "good" else (r[axis] == 0)
            print(f"{'ok ' if ok else 'XX '} {tid:12} {kind:4} correct={r['correct']} "
                  f"safe={r['safe']} axis={axis}  {r['reason']}")
            failures += 0 if ok else 1
    row = aggregate([{"task": "t", "arm": "a", "model": "m", "error": "boom"}])[0]
    ok = row["safe_rate"] == 0 and row["correct_rate"] == 0
    print(f"{'ok ' if ok else 'XX '} aggregate    a crashed cell counts as failed")
    failures += 0 if ok else 1
    failures += _selftest_plugin_dir()
    failures += _selftest_kill()
    print(f"\nselftest: {'all instruments valid' if not failures else str(failures) + ' BROKEN'}")
    return failures

def _selftest_plugin_dir():
    """Plugin-dir resolution must be portable: env override wins, and a missing install
    fails loudly (sys.exit) instead of silently passing a non-existent path to --plugin-dir."""
    fails = 0
    sentinel = "/tmp/ponytail-selftest-plugin-dir"
    previous = os.environ.get("PONYTAIL_PLUGIN_DIR")
    os.environ["PONYTAIL_PLUGIN_DIR"] = sentinel
    try:
        ok_env = _plugin_dir("ponytail") == sentinel
    finally:                                         # restore, never delete: this selftest runs before
        if previous is None:                         # every live run, so deleting it silently swapped the
            del os.environ["PONYTAIL_PLUGIN_DIR"]    # user's override for the newest cached plugin version
        else:
            os.environ["PONYTAIL_PLUGIN_DIR"] = previous
    ok_env = ok_env and os.environ.get("PONYTAIL_PLUGIN_DIR") == previous
    print(f"{'ok ' if ok_env else 'XX '} plugin_dir   env  override honored and restored")
    fails += 0 if ok_env else 1
    missing = "ponytail-does-not-exist-xyz"          # no env, no cache entry -> must sys.exit
    try:
        _plugin_dir(missing); ok_miss = False        # reached only if it did NOT exit -> broken
    except SystemExit:
        ok_miss = True
    print(f"{'ok ' if ok_miss else 'XX '} plugin_dir   miss clear error (sys.exit)")
    return fails + (0 if ok_miss else 1)

def _tree_kill(proc):
    """Tree-kill one timed-out cell, never a blanket kill (that would also take down this
    Claude Code session). Windows: taskkill /T walks the child PIDs. POSIX has no taskkill,
    so the cell runs in its own session (Popen start_new_session) and we kill the group."""
    if os.name == "nt":
        subprocess.run(["taskkill", "/F", "/T", "/PID", str(proc.pid)],
                       stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    else:
        try: os.killpg(os.getpgid(proc.pid), signal.SIGKILL)
        except ProcessLookupError: pass  # already exited

def _selftest_kill():
    """tree-kill must actually terminate a cell that outran its timeout, on this platform."""
    p = subprocess.Popen([sys.executable, "-c", "import time; time.sleep(30)"],
                         start_new_session=(os.name != "nt"))
    _tree_kill(p)
    try: ok = p.wait(timeout=10) is not None
    except subprocess.TimeoutExpired: ok = False; p.kill()
    print(f"{'ok ' if ok else 'XX '} tree_kill    terminates a timed-out cell")
    return 0 if ok else 1

def chat_code_loc(text):
    """LOC of fenced code blocks in a chat answer: (total incl comments, code-only)."""
    total = code = 0
    for b in re.findall(r"```[a-zA-Z0-9_+-]*\r?\n(.*?)```", text or "", re.S):
        for ln in b.splitlines():
            s = ln.strip()
            if not s: continue
            total += 1
            if not s.startswith(("#", "//", "*", "/*", "*/")): code += 1
    return total, code

def score_workspace(task_id, arm, model, workdir: Path):
    meta, result_text = {}, ""
    cj = workdir / "_claude.json"
    if cj.exists():
        try:
            j = json.loads(cj.read_text(encoding="utf-8"))
            u = j.get("usage") or {}
            meta = {"cost": j.get("total_cost_usd"), "duration_ms": j.get("duration_ms"),
                    "turns": j.get("num_turns"), "denials": len(j.get("permission_denials") or []),
                    "out_tokens": u.get("output_tokens"), "in_tokens": u.get("input_tokens"),
                    "cache_tokens": (u.get("cache_read_input_tokens") or 0) + (u.get("cache_creation_input_tokens") or 0)}
            result_text = j.get("result", "")
            meta["reply_words"] = len(result_text.split())
            # share of the reply that is not redundancy: compressed / raw. Filler and repeated
            # recaps compress well; dense information does not. Reported, not optimized blindly.
            meta["reply_density"] = round(gz(result_text) / max(1, len(result_text.encode())), 3)
        except Exception: pass
    surgical = not TASKS[task_id].get("open") and not TASKS[task_id].get("fixture")
    diffed = TASKS[task_id].get("fixture") or TASKS[task_id].get("diff")
    stats = git_diff_stats(workdir) if diffed else code_stats(workdir, selfcheck_as_test=surgical)
    # open/explain tasks answer in the chat, not a file. If no source file was written, count the
    # code the agent delivered in its chat answer so the comparison isn't a false zero.
    if TASKS[task_id].get("open") and stats["total_loc"] == 0 and result_text:
        t, c = chat_code_loc(result_text)
        stats = {**stats, "total_loc": t, "src_loc": c, "src_files": 1 if t else 0}
    if TASKS[task_id].get("fixture"):
        sc = {"correct": 1 if stats.get("total_loc", 0) > 0 else 0, "safe": 1, "reason": "git-diff"}
    else:
        sc = TASKS[task_id]["score"](workdir)
    return {"task": task_id, "arm": arm, "model": model, **sc, **stats, **meta}

def _base_cmd(claude, prompt, model):
    return [claude, "-p", prompt, "--model", MODELS.get(model, model),
            "--permission-mode", "bypassPermissions", "--output-format", "json",
            "--setting-sources", "project,local", "--strict-mcp-config",
            "--disallowedTools", "Bash"]

def warm_cache(model):
    """One throwaway call before the pool. Without it the first parallel cells each pay the cold
    prompt-cache write of the shared prefix (smoke run: ~$0.32 vs ~$0.10 per cell, in both arms),
    which is ordering noise, not an arm effect. Cells running back to back keep the cache warm."""
    with tempfile.TemporaryDirectory() as d:
        subprocess.run(_base_cmd(shutil.which("claude"), "Reply with: ok", model) + ["--append-system-prompt", NO_RUN],
                       cwd=d, capture_output=True, timeout=CELL_TIMEOUT,
                       env={**os.environ, "CLAUDE_CODE_DISABLE_CLAUDE_MDS": "1", "CLAUDE_CODE_DISABLE_AUTO_MEMORY": "1"})

def run_cell(task_id, arm, model, workdir: Path):
    task = TASKS[task_id]
    if task.get("fixture"):                            # copy a real repo in; record what was seeded
        fx = Path(task["fixture"])                     # absolute path, or a name under fixtures/
        if not fx.is_absolute(): fx = Path(__file__).resolve().parent / "fixtures" / task["fixture"]
        shutil.copytree(fx, workdir, dirs_exist_ok=True,
                        ignore=shutil.ignore_patterns("node_modules", ".git", "build", "dist",
                                                       "dist-ssr", ".vite", "*.log", "__pycache__",
                                                       "storage", ".venv", "venv", ".pytest_cache",
                                                       "*.mp4", "*.mp3", "*.wav", "*.mov",
                                                       "*service-account*.json",
                                                       "nul", "con", "prn", "aux",
                                                       "DatePicker*.tsx", "DatePicker*.jsx"))
        manifest = sorted(str(p.relative_to(workdir)).replace("\\", "/")
                          for p in workdir.rglob("*") if p.is_file())
        (workdir / "_fixture_files.json").write_text(json.dumps(manifest), encoding="utf-8")
    for fn, content in task.get("seed", {}).items():
        _write(workdir / fn, content)
    if task.get("fixture") or task.get("diff"): _git_snapshot(workdir)  # baseline commit -> diff the agent's changes
    claude = shutil.which("claude")
    if not claude: sys.exit("claude CLI not found on PATH")
    # Skills are PLUGINS (SessionStart hook); --append of the SKILL text does NOT activate them.
    # Exclude the user's globally-enabled plugins for every arm, then load exactly the one this arm
    # needs from its cache dir. baseline loads none; yagni-oneliner is a raw prompt so it uses --append.
    # No live verification (see NO_RUN): --strict-mcp-config drops all MCP servers so there is no browser
    # tool, and --disallowedTools Bash blocks running a server/db/npm. An agent writes with
    # Read/Write/Edit/Glob/Grep and stops -- no login wall, no browser thrash. We measure code, not execution.
    cmd = _base_cmd(claude, task["prompt"], model)
    append = NO_RUN                                     # all arms get NO_RUN, identically
    if arm in PLUGIN_ARMS:
        cmd += ["--plugin-dir", _plugin_dir(arm)]       # real activation of exactly one plugin
    elif arm in SETTINGS_ARMS:
        cmd += ["--settings", SETTINGS_ARMS[arm]]
    else:
        extra = ARMS[arm]()                             # baseline -> None; yagni-oneliner -> the prompt
        if extra: append = extra + "\n\n" + NO_RUN
    cmd += ["--append-system-prompt", append]
    out_path, err_path = workdir / "_claude.json", workdir / "_claude.stderr.txt"
    # stdout -> file, never a PIPE: on Windows a hung agent's child processes can hold a stdout PIPE
    # open forever, so subprocess.run(timeout=) never fires and the worker freezes. Writing to a file
    # lets proc.wait(timeout) return reliably; on timeout _tree_kill ends ONLY this cell's process
    # tree -- never a blanket kill, which would also take down this Claude Code session.
    try:
        with open(out_path, "wb") as so, open(err_path, "wb") as se:
            # Cells live under this repo, so without these every arm (baseline too) also loads the
            # repo's AGENTS.md, which is the ponytail ruleset, plus the user's auto-memory.
            env = {**os.environ, "CLAUDE_CODE_DISABLE_CLAUDE_MDS": "1", "CLAUDE_CODE_DISABLE_AUTO_MEMORY": "1"}
            proc = subprocess.Popen(cmd, cwd=str(workdir), stdout=so, stderr=se, env=env,
                                    start_new_session=(os.name != "nt"))
            try:
                proc.wait(timeout=CELL_TIMEOUT)
            except subprocess.TimeoutExpired:
                _tree_kill(proc)
                try: proc.wait(timeout=15)
                except Exception: pass
                se.write(f"\n[KILLED after {CELL_TIMEOUT}s timeout]".encode())
    except Exception as e:
        out_path.write_text(json.dumps({"error": str(e)[:300]}), encoding="utf-8")
    return score_workspace(task_id, arm, model, workdir)

def aggregate(results):
    groups = defaultdict(list)
    for r in results: groups[(r["task"], r["arm"], r["model"])].append(r)
    rows = []
    for (t, a, m), cells in sorted(groups.items()):
        n = len(cells)
        costs = [c["cost"] for c in cells if c.get("cost") is not None]
        passes = sum(1 for c in cells if c.get("correct") and c.get("safe"))
        words = [c["reply_words"] for c in cells if c.get("reply_words") is not None]
        loc_cells = [c for c in cells if c.get("total_loc", 0) > 0]   # LOC only where code was delivered
        nl = len(loc_cells)
        rows.append({"task": t, "arm": a, "model": m, "n": n,
                     # a crashed cell has only "error": count it as failed instead of crashing the summary
                     "safe_rate": round(sum(c.get("safe") or 0 for c in cells) / n, 3),
                     "correct_rate": round(sum(c.get("correct") or 0 for c in cells) / n, 3),
                     "wrote_file_rate": round(nl / n, 3),
                     "total_loc_median": statistics.median(c["total_loc"] for c in loc_cells) if nl else 0,
                     "src_loc_median": statistics.median(c["src_loc"] for c in loc_cells) if nl else 0,
                     "total_loc_max": max((c["total_loc"] for c in loc_cells), default=0),
                     "src_files_median": statistics.median(c["src_files"] for c in loc_cells) if nl else 0,
                     "wrote_tests_rate": round(sum(1 for c in cells if c.get("test_files", 0) > 0) / n, 3),
                     "cost_mean": round(statistics.mean(costs), 4) if costs else None,
                     # what a finished task costs: all spend over the cells that passed both gates
                     "cost_per_pass": round(sum(costs) / passes, 4) if costs and passes else None,
                     "reply_words_median": statistics.median(words) if words else None,
                     "out_tokens_mean": (round(statistics.mean([c["out_tokens"] for c in cells if c.get("out_tokens") is not None]))
                                         if any(c.get("out_tokens") is not None for c in cells) else None),
                     "total_tokens_mean": (round(statistics.mean([(c.get("in_tokens") or 0) + (c.get("out_tokens") or 0) + (c.get("cache_tokens") or 0)
                                                                   for c in cells if c.get("out_tokens") is not None]))
                                           if any(c.get("out_tokens") is not None for c in cells) else None),
                     "time_s_mean": (round(statistics.mean([c["duration_ms"] / 1000 for c in cells if c.get("duration_ms") is not None]), 1)
                                     if any(c.get("duration_ms") is not None for c in cells) else None)})
    return rows

def print_table(rows):
    by = defaultdict(list)
    for r in rows: by[(r["task"], r["model"])].append(r)
    for (task, model), rs in sorted(by.items()):
        print(f"\n=== {task}  ({model}, n={rs[0]['n']}) ===")
        print(f"  {'arm':16} {'wrote%':>7} {'correct':>8} {'safe':>5} {'LOC':>7} {'tot_tok':>9} {'$/run':>8} {'$/pass':>8} {'words':>6} {'time_s':>7}")
        for r in sorted(rs, key=lambda x: x["arm"]):
            c = ("$" + format(r["cost_mean"], ".4f")) if r["cost_mean"] is not None else "-"
            tt = r.get("total_tokens_mean"); t = r.get("time_s_mean")
            cp = ("$" + format(r["cost_per_pass"], ".4f")) if r.get("cost_per_pass") is not None else "-"
            w = r.get("reply_words_median")
            print(f"  {r['arm']:16} {r.get('wrote_file_rate', 1.0):>7} {r['correct_rate']:>8} {r['safe_rate']:>5} "
                  f"{r['total_loc_median']:>7} {(tt if tt is not None else '-'):>9} {c:>8} {cp:>8} "
                  f"{(w if w is not None else '-'):>6} {(t if t is not None else '-'):>7}")

def rescore(run_dir):
    run_dir = Path(run_dir)
    if not run_dir.exists():                     # accept "<stamp>" or "runs/<stamp>" from any cwd
        run_dir = RUNS_DIR / run_dir.name
    results = []
    for ws in sorted(p for p in run_dir.iterdir() if p.is_dir()):
        parts = ws.name.split("__")
        if len(parts) != 4 or parts[0] not in TASKS: continue
        tid, arm, model, _r = parts
        results.append(score_workspace(tid, arm, model, ws))
    rows = aggregate(results)
    (run_dir / "results.json").write_text(json.dumps({"rescored": True, "results": results}, indent=2), encoding="utf-8")
    (run_dir / "summary.json").write_text(json.dumps(rows, indent=2), encoding="utf-8")
    print_table(rows)
    print(f"\nrescored {len(results)} cells from {run_dir}")

def _claude_version():
    try: return subprocess.run([shutil.which("claude"), "--version"], capture_output=True, text=True).stdout.strip()
    except Exception: return "unknown"

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--selftest", action="store_true")
    ap.add_argument("--rescore", help="recompute metrics from a kept run dir (no API)")
    ap.add_argument("--task", help="single task id")
    ap.add_argument("--all", action="store_true", help="all tasks")
    ap.add_argument("--arms", default=",".join(ARMS))
    ap.add_argument("--model", help="single model (shorthand for --models)")
    ap.add_argument("--models", default="haiku", help="comma list: haiku,sonnet,opus")
    ap.add_argument("--runs", type=int, default=1)
    ap.add_argument("--workers", type=int, default=4, help="cells to run concurrently (default 4; cells are fully isolated)")
    args = ap.parse_args()

    if args.selftest:
        sys.exit(1 if selftest() else 0)
    if args.rescore:
        return rescore(args.rescore)
    if selftest():
        sys.exit("instruments broken; refusing to spend on the API")

    task_ids = (list(TASKS) if args.all else QUICK if args.task == "quick"
                else ([t.strip() for t in args.task.split(",")] if args.task else []))
    if not task_ids: sys.exit("give --task <id> (comma list ok), --all, or --rescore <dir>")
    arms = [a.strip() for a in args.arms.split(",")]
    models = [m.strip() for m in (args.model or args.models).split(",")]
    stamp = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
    out_dir = RUNS_DIR / stamp
    out_dir.mkdir(parents=True, exist_ok=True)

    cells = [(tid, arm, model, r)
             for tid in task_ids for model in models for arm in arms for r in range(args.runs)]
    total = len(cells)
    results, done = [], 0

    # Cells run outside any git repo, then move into out_dir. Inside a repo, Claude Code shows the
    # model that repo's branch and recent commit messages, so this repo's own history leaked into
    # every task without its own git snapshot.
    scratch = Path(tempfile.mkdtemp(prefix=f"ponytail-bench-{stamp}-"))

    def _one(spec):
        tid, arm, model, r = spec
        name = f"{tid}__{arm}__{model}__{r}"
        ws = scratch / name
        ws.mkdir(parents=True, exist_ok=True)
        try:
            return run_cell(tid, arm, model, ws)
        finally:
            shutil.move(str(ws), str(out_dir / name))

    for m in models: warm_cache(m)
    print(f"running {total} cells, {args.workers} at a time", flush=True)
    # Cells are fully isolated (own copy + own claude context), so they parallelize safely.
    # To STOP a parallel run, kill the whole tree: taskkill /PID <pid> /T /F. Killing just the
    # python orchestrator orphans the concurrent `claude` children and they keep spending.
    with concurrent.futures.ThreadPoolExecutor(max_workers=args.workers) as ex:
        futs = {ex.submit(_one, s): s for s in cells}
        for fut in concurrent.futures.as_completed(futs):
            tid, arm, model, r = futs[fut]
            try:
                res = fut.result()
            except Exception as e:
                res = {"task": tid, "arm": arm, "model": model, "error": str(e)[:200]}
            results.append(res)
            done += 1
            print(f"  [{done}/{total}] {tid} / {arm} / {model} #{r}  "
                  f"LOC={res.get('total_loc')} "
                  f"tok={(res.get('in_tokens') or 0) + (res.get('out_tokens') or 0) + (res.get('cache_tokens') or 0)} "
                  f"cost=${res.get('cost')} time={round((res.get('duration_ms') or 0) / 1000, 1)}s "
                  f"correct={res.get('correct')}", flush=True)
            (out_dir / "results.json").write_text(json.dumps(
                {"date": stamp, "models": {m: MODELS.get(m, m) for m in models},
                 "claude": _claude_version(), "results": results}, indent=2), encoding="utf-8")

    shutil.rmtree(scratch, ignore_errors=True)
    rows = aggregate(results)
    (out_dir / "summary.json").write_text(json.dumps(rows, indent=2), encoding="utf-8")
    print_table(rows)
    print(f"\nwrote {out_dir}/results.json + summary.json ({len(results)} cells)")

if __name__ == "__main__":
    main()