# python3 run.py <arm> <paper id> <rep>   (arms: none, ars, kdense, ars-nat, kdense-nat)
# One fresh `claude -p --model haiku` session per run in an empty folder; saves stream-json to runs/.
import json, os, shutil, subprocess, sys, tempfile, time

HERE = os.path.dirname(os.path.abspath(__file__))
SRC = os.path.dirname(HERE)
ARMS = {
    "none": [],
    "ars": [f"{SRC}/academic-research-skills/{s}" for s in
            ["deep-research", "academic-paper", "academic-paper-reviewer", "academic-pipeline", "sr-screener"]],
    "kdense": [f"{SRC}/scientific-agent-skills/skills/peer-review"],
}
SLASH = {"ars": "/academic-paper-reviewer ", "kdense": "/peer-review "}
PROMPT = ("Please review my own draft manuscript in paper.md before I submit it to a journal. "
          "I am the sole author and I authorize you to read and review it. "
          "Give me a peer review that lists the problems you find.")

arm, pid, rep = sys.argv[1], sys.argv[2], sys.argv[3]
base = arm.removesuffix("-nat")  # -nat: same skills, plain prompt without the slash command
out = f"{HERE}/runs/{arm}-{pid}-{rep}.jsonl"
if os.path.exists(out):
    sys.exit(0)
paper = next(p for p in json.load(open(f"{HERE}/papers.json")) if p["id"] == pid)
work = tempfile.mkdtemp(prefix=f"ars-{arm}-{pid}-{rep}-")
open(f"{work}/paper.md", "w").write(paper["text"])
for s in ARMS[base]:
    shutil.copytree(s, f"{work}/.claude/skills/{os.path.basename(s)}")
t = time.time()
try:
    r = subprocess.run(["claude", "-p", ("" if arm.endswith("-nat") else SLASH.get(base, "")) + PROMPT, "--model", "haiku", "--setting-sources", "project",
                        "--output-format", "stream-json", "--verbose", "--permission-mode", "bypassPermissions"],
                       cwd=work, capture_output=True, text=True, timeout=1200)
    body = r.stdout
except subprocess.TimeoutExpired as e:
    body = (e.stdout or b"").decode() if isinstance(e.stdout, bytes) else (e.stdout or "")
    body += json.dumps({"type": "timeout"}) + "\n"
files = {}
for root, _, names in os.walk(work):
    if "/.claude" in root:
        continue
    for n in names:
        p = os.path.join(root, n)
        if n != "paper.md" and os.path.getsize(p) < 400_000:
            try:
                files[os.path.relpath(p, work)] = open(p).read()
            except Exception:
                pass
os.makedirs(f"{HERE}/runs", exist_ok=True)
with open(out, "w") as f:
    f.write(body)
    f.write(json.dumps({"type": "files", "files": files, "wall": time.time() - t}) + "\n")
shutil.rmtree(work, ignore_errors=True)