#!/usr/bin/env python3
"""collect.py — planted-numbers trace test: collect one answer per (question, assistant).

Rules: PREREG — planted-numbers — trace test — 2026-11-08 (hash-locked before this runs for real).
Subjects: Claude Sonnet via `claude -p` (WebSearch, WebFetch); GPT-5.6 Sol via Codex (web_search).
Each runs from a fresh, neutrally named empty temp dir outside the vault. No retries on content, only on
empty/failed calls (max 3). Output: answers.jsonl (one row per call), never edited by hand.

    python3 collect.py            # all 40 calls
    python3 collect.py --dry      # print the jobs, call nothing
"""
import json, os, subprocess, sys, tempfile, time, hashlib
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path

HERE = Path(__file__).resolve().parent
Q = json.loads((HERE / "questions.json").read_text())
OUT = HERE / "answers.jsonl"
SYSTEM = "You are a helpful assistant."
SCRATCH = Path(os.environ.get("PN_SCRATCH", tempfile.gettempdir()))
CODEX = next((c for c in ("/Applications/ChatGPT.app/Contents/Resources/codex-cli/bin/codex",
                           "/Applications/ChatGPT.app/Contents/Resources/codex") if os.path.exists(c)), "codex")
GPT_MODEL = "gpt-5.6-sol"


def fresh_dir():
    return Path(tempfile.mkdtemp(prefix="t-", dir=SCRATCH))


def call_claude(msg):
    cwd = fresh_dir()
    cmd = ["claude", "-p", msg, "--system-prompt", SYSTEM, "--setting-sources", "",
           "--strict-mcp-config", "--mcp-config", '{"mcpServers":{}}',
           "--model", "sonnet", "--no-session-persistence", "--output-format", "json",
           "--tools", "WebSearch,WebFetch", "--allowedTools", "WebSearch", "WebFetch"]
    last = ""
    for attempt in range(3):
        p = subprocess.run(cmd, cwd=cwd, capture_output=True, text=True, timeout=900, stdin=subprocess.DEVNULL)
        raw = p.stdout.strip()
        try:
            d = json.loads(raw[raw.index("{"):])
        except Exception:
            last = p.stderr[-400:]; time.sleep(20 * (attempt + 1)); continue
        text = d.get("result") or ""
        if text.strip() and not d.get("is_error"):
            return {"text": text, "model_version": ",".join((d.get("modelUsage") or {}).keys()),
                    "num_turns": d.get("num_turns"), "cwd": cwd.name, "attempts": attempt + 1}
        last = raw[-400:]; time.sleep(20 * (attempt + 1))
    raise RuntimeError(f"claude failed 3x: {last}")


def call_gpt(msg):
    cwd = fresh_dir()
    fd, outf = tempfile.mkstemp(suffix=".txt", dir=cwd); os.close(fd)
    cmd = [CODEX, "exec", "--ignore-user-config", "-m", GPT_MODEL, "-s", "read-only", "--skip-git-repo-check",
           "--ephemeral", "-C", str(cwd), "--json", "-o", outf, "-c", "tools.web_search=true", "-"]
    p = None
    for attempt in range(3):
        p = subprocess.run(cmd, input=msg, capture_output=True, text=True, timeout=900)
        text = Path(outf).read_text().strip() if Path(outf).exists() else ""
        if text:
            searches = sum(1 for l in p.stdout.splitlines() if '"type":"web_search"' in l and '"item.completed"' in l)
            Path(outf).unlink(missing_ok=True)
            return {"text": text, "model_version": GPT_MODEL, "search_queries": searches, "cwd": cwd.name,
                    "attempts": attempt + 1}
        time.sleep(20 * (attempt + 1))
    raise RuntimeError(f"codex empty 3x: {p.stderr[-400:] if p else ''}")


SUBJECTS = {"claude": call_claude, "gpt": call_gpt}


def job(q, subject):
    msg = Q["template"].format(topic=q["topic"])
    t0 = time.time()
    try:
        r = SUBJECTS[subject](msg)
        err = None
    except Exception as e:
        r, err = {"text": ""}, str(e)[:500]
    row = {"qid": q["id"], "subject": subject, "prompt": msg, "error": err,
           "started_utc": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime(t0)),
           "secs": round(time.time() - t0, 1), **r}
    row["text_sha256"] = hashlib.sha256(row["text"].encode()).hexdigest()
    return row


def main():
    jobs = [(q, s) for q in Q["questions"] for s in SUBJECTS]
    if "--dry" in sys.argv:
        for q, s in jobs:
            print(s, q["id"], Q["template"].format(topic=q["topic"]))
        return
    if OUT.exists():
        sys.exit(f"{OUT.name} exists: refusing to overwrite a collected run")
    pools = {"claude": ThreadPoolExecutor(3), "gpt": ThreadPoolExecutor(3)}
    futs = [pools[s].submit(job, q, s) for q, s in jobs]
    with OUT.open("w") as f:
        for fu in futs:
            row = fu.result()
            f.write(json.dumps(row, ensure_ascii=False) + "\n"); f.flush()
            print(row["subject"], row["qid"], "ERR" if row["error"] else "ok", row["secs"], flush=True)


if __name__ == "__main__":
    main()