"""Arm A of the head-to-head: plain Claude Code, one session, no throng.

Runs a headless Claude Code session (`claude -p`, the user's own login) in the
repository for a fixed wall-clock time. If the session ends early it is resumed
with "keep going" until the time is up. The agent works the normal Claude Code
way (its own tools, `cargo test`, its own subagents); throng is not involved in
its loop.

Measurement is done from outside and the agent never sees it: every
CHECKPOINT_MIN minutes and at the end, the untouched scaffold plus the agent's
current `src/` is graded with the same grader as Arm B, and committed with the
score (so editing tests or the grader cannot raise the score; any change
outside src/ is recorded). Every session's full output is saved in
proof/runs/solo/.

    python scaffold/arm_solo.py --repo <path-to>/sqlswarm-solo --minutes 30 --model opus --goal-file scaffold/goal.txt
"""

import argparse
import json
import os
import shutil
import subprocess
import tempfile
import threading
import time
from pathlib import Path

CHECKPOINT_MIN = 5
NO_WINDOW = 0x08000000
# The Claude Code executable (npm's `claude` on PATH is a .cmd wrapper).
CLAUDE = os.environ.get("THRONG_CLAUDE") or shutil.which("claude.exe") or str(
    Path(os.environ.get("APPDATA", "")) / "npm" / "node_modules" / "@anthropic-ai" / "claude-code" / "bin" / "claude.exe")


def sh(args, cwd, **kw):
    return subprocess.run(args, cwd=cwd, capture_output=True, text=True, encoding="utf-8", errors="replace", creationflags=NO_WINDOW, **kw)


def grade(repo: Path, work: Path, n: int):
    """Grade scaffold + current src/ (a sanitized copy), return (passed, total, touched, progress detail)."""
    touched = [l[3:] for l in sh(["git", "status", "--porcelain"], repo).stdout.splitlines() if l[3:] and not l[3:].startswith(("src/", "proof/", "target/"))]
    copy = Path(tempfile.mkdtemp(prefix="solo-grade-"))
    try:
        arch = subprocess.run(["git", "archive", "scaffold"], cwd=repo, capture_output=True, creationflags=NO_WINDOW).stdout
        subprocess.run(["tar", "-x", "-C", str(copy)], input=arch, creationflags=NO_WINDOW)
        shutil.rmtree(copy / "src", ignore_errors=True)
        shutil.copytree(repo / "src", copy / "src")
        prev = work / f"grade-{n - 1}.json"
        cmd = ["throng", "grade", "--repo", str(copy), "--json", "--agent", "solo-grader", "--save", str(work / f"grade-{n}.json")]
        if prev.exists():
            cmd += ["--baseline", str(prev)]
        r = sh(cmd, repo)
        v = json.loads(r.stdout) if r.stdout.strip().startswith("{") else {}
        detail = (v.get("progress") or {}).get("detail", "")
        return v.get("passed", 0), v.get("total", 0), touched, detail
    finally:
        shutil.rmtree(copy, ignore_errors=True)


def checkpoint(repo: Path, work: Path, n: int, t0: float, model: str, sessions: list, final=False):
    passed, total, touched, detail = grade(repo, work, n)
    mins = (time.time() - t0) / 60
    entry = {"ts_ms": int(time.time() * 1000), "arm": "solo", "checkpoint": n, "final": final, "elapsed_min": round(mins, 1), "model": model,
             "tests_after": passed, "tests_total": total, "vs_previous": detail, "touched_read_only": touched, "sessions": sessions}
    ledger = repo / "proof" / "ledger.jsonl"
    with open(ledger, "a", encoding="utf-8") as f:
        f.write(json.dumps(entry) + "\n")
    sh(["git", "add", "--", "src", "proof"], repo)
    msg = (f"{'final' if final else 'checkpoint'} at {mins:.0f} min: {passed} of {total} tests pass\n\n"
           f"Arm: solo (plain Claude Code)\nModel: {model}\nElapsed-Min: {mins:.1f}\nTests-After: {passed}\nTests-Total: {total}\n"
           f"Vs-Previous: {detail}\nTouched-Read-Only: {', '.join(touched) or 'none'}\n")
    sh(["git", "-c", "user.name=claude code (solo arm)", "-c", "user.email=solo@sqlswarm.local", "commit", "-q", "--allow-empty", "-m", msg], repo)
    print(f"[{mins:5.1f} min] checkpoint {n}: {passed}/{total} tests pass  {detail}", flush=True)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--repo", required=True)
    ap.add_argument("--minutes", type=float, default=30)
    ap.add_argument("--model", default="opus")
    ap.add_argument("--goal-file", required=True)
    a = ap.parse_args()
    repo = Path(a.repo).resolve()
    runs = repo / "proof" / "runs" / "solo"
    runs.mkdir(parents=True, exist_ok=True)
    work = Path(tempfile.mkdtemp(prefix="solo-work-"))
    goal = (repo / a.goal_file).read_text(encoding="utf-8").strip()
    first = (f"{goal}\n\nOnly change files under src/; everything else (the grader in slt/, tests/, corpus/) is read-only. "
             "Run `cargo test` to see how many tests pass (each test grades 10 sqllogictest records). "
             "Work autonomously and keep going until every test passes; do not stop to ask questions. Do not use the internet.")
    deadline = time.time() + a.minutes * 60
    t0 = time.time()
    sessions, n, session_id = [], 0, None
    stop = threading.Event()

    def ticker():
        nonlocal n
        while not stop.wait(CHECKPOINT_MIN * 60):
            if time.time() < deadline:
                n += 1
                checkpoint(repo, work, n, t0, a.model, list(sessions))

    threading.Thread(target=ticker, daemon=True).start()
    part = 0
    while time.time() < deadline:
        part += 1
        prompt = first if session_id is None else "Keep going: make more of the tests pass."
        cmd = [CLAUDE, "-p", "--output-format", "stream-json", "--verbose", "--model", a.model,
               "--permission-mode", "acceptEdits", "--allowedTools", "Bash", "Read", "Write", "Edit", "Glob", "Grep", "Task",
               "--disallowedTools", "WebSearch", "WebFetch", "--strict-mcp-config", "--disable-slash-commands",
               "--settings", '{"disableAllHooks": true}']
        if session_id:
            cmd += ["--resume", session_id]
        cmd.append(prompt)
        out = open(runs / f"session-part{part}.jsonl", "w", encoding="utf-8")
        p = subprocess.Popen(cmd, cwd=repo, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, stdin=subprocess.DEVNULL, text=True, encoding="utf-8", errors="replace", creationflags=NO_WINDOW)
        killer = threading.Timer(max(1, deadline - time.time()), lambda: subprocess.run(["taskkill", "/T", "/F", "/PID", str(p.pid)], capture_output=True, creationflags=NO_WINDOW))
        killer.start()
        for line in p.stdout:
            out.write(line)
            try:
                v = json.loads(line)
            except ValueError:
                continue
            if v.get("session_id"):
                session_id = v["session_id"]
            if v.get("type") == "result":
                sessions.append({"part": part, "session": session_id, "cost_usd_equiv": v.get("total_cost_usd"), "turns": v.get("num_turns"), "subtype": v.get("subtype")})
        p.wait()
        killer.cancel()
        out.close()
        print(f"[{(time.time() - t0) / 60:5.1f} min] session part {part} ended", flush=True)
        if session_id is None:
            break
    stop.set()
    n += 1
    checkpoint(repo, work, n, t0, a.model, sessions, final=True)


if __name__ == "__main__":
    main()