#!/usr/bin/env python3
"""Run a solver over a SWE-bench Verified subset with the real Docker eval harness.

Pipeline per instance: clone the repo at base_commit, run the solver on the
problem statement to produce a source patch (tests restored so it can't edit the
grader's tests), capture cost, then evaluate every patch through the official
`swebench` harness (build image, apply patch + hidden test patch, run
FAIL_TO_PASS + PASS_TO_PASS). Reports resolution and dollars per resolved task.

With --rollouts N, runs N independent rollouts per instance and reports the
**metered test-time-compute frontier**: best-of-k resolution (an instance counts
resolved if any of the first k rollouts resolves it) paired with the cumulative
dollars those k rollouts cost — the resolution × $ curve of test-time compute,
metered. (Literature: best-of-N + selection lifts frontier agents ~7pp on
SWE-bench Verified; this measures what that lift costs.)

Setup (macOS, Apple Silicon):
    open -a OrbStack                       # Docker daemon (emulates linux/amd64)
    python3.11 -m venv ~/.swebench-venv    # swebench needs Python 3.10+
    ~/.swebench-venv/bin/pip install swebench datasets
    ~/.swebench-venv/bin/python swebench/run_swebench.py --solver claude \
        --instances psf__requests-5414 pylint-dev__pylint-7080 --rollouts 3
"""
import argparse, json, subprocess, sys, shutil
from pathlib import Path

def sh(cmd, cwd=None, timeout=1800):
    return subprocess.run(cmd, shell=True, cwd=cwd, capture_output=True, text=True, timeout=timeout)

def claude_solve(repo_dir, problem):
    prompt = (
        f"You are fixing a real bug in this checkout. Edit ONLY the library source "
        f"(NOT tests — the grader applies its own). Explore the code, make the fix, "
        f"then run the repo's existing pytest for the affected module to confirm you "
        f"did not break anything; iterate until those pass.\n\nISSUE:\n{problem}"
    )
    safe = prompt.replace("'", "'\\''")
    out = sh(f"claude -p '{safe}' --dangerously-skip-permissions --output-format json", cwd=repo_dir)
    try:
        return json.loads(out.stdout).get("total_cost_usd", 0.0) or 0.0
    except Exception:
        return 0.0

def one_rollout(ds, instances, solver, work, dataset, run_id):
    """Generate a patch per instance and evaluate the batch. Returns (resolved_set, total_cost)."""
    preds, total_cost = [], 0.0
    for iid in instances:
        inst = ds[iid]
        repo_dir = work / f"{iid}__{run_id}"
        if repo_dir.exists():
            shutil.rmtree(repo_dir)
        sh(f"git clone -q https://github.com/{inst['repo']} {repo_dir}", timeout=600)
        sh(f"git checkout -q {inst['base_commit']}", cwd=repo_dir)
        cost = claude_solve(repo_dir, inst["problem_statement"]) if solver == "claude" else 0.0
        total_cost += cost
        sh("git checkout -- $(git diff --name-only | grep -E '(^|/)(test_|tests/)' || true)", cwd=repo_dir)
        patch = sh("git diff", cwd=repo_dir).stdout
        preds.append({"instance_id": iid, "model_name_or_path": solver, "model_patch": patch})
        print(f"  [{iid}] cost=${cost:.3f} patch_bytes={len(patch)}", flush=True)
    preds_path = work / f"predictions_{run_id}.json"
    json.dump(preds, open(preds_path, "w"))
    subprocess.run(
        [sys.executable, "-m", "swebench.harness.run_evaluation",
         "--dataset_name", dataset, "--predictions_path", str(preds_path),
         "--max_workers", "2", "--run_id", run_id,
         "--instance_ids", *instances, "--cache_level", "instance"],
        check=False, capture_output=True)
    report = Path(f"{solver}.{run_id}.json")
    resolved = set(json.load(open(report))["resolved_ids"]) if report.exists() else set()
    return resolved, total_cost

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--solver", default="claude", choices=["claude"])
    ap.add_argument("--instances", nargs="+", required=True)
    ap.add_argument("--dataset", default="princeton-nlp/SWE-bench_Verified")
    ap.add_argument("--run-id", default="joule-bench-swe")
    ap.add_argument("--rollouts", type=int, default=1)
    ap.add_argument("--work", default=str(Path.home() / "data-share" / "swebench-run"))
    args = ap.parse_args()

    from datasets import load_dataset
    ds = {r["instance_id"]: r for r in load_dataset(args.dataset, split="test")}
    work = Path(args.work); work.mkdir(parents=True, exist_ok=True)
    n = len(args.instances)

    per_rollout = []  # (resolved_set, cost) per rollout
    for k in range(1, args.rollouts + 1):
        print(f"\n=== rollout {k}/{args.rollouts} ===", flush=True)
        per_rollout.append(one_rollout(ds, args.instances, args.solver, work, args.dataset, f"{args.run_id}-r{k}"))

    # Best-of-k metered test-time-compute frontier.
    print("\n" + "=" * 60)
    if args.rollouts == 1:
        resolved, cost = per_rollout[0]
        res = len(resolved); per = f"${cost/res:.2f}" if res else "—"
        print(f"  RESOLVED {res}/{n} = {100*res/n:.0f}%   total ${cost:.2f}   $/resolved {per}")
    else:
        print("  test-time-compute frontier (best-of-k):")
        print(f"  {'k':>2} {'resolved':>10} {'resol.%':>8} {'cum $':>9}")
        union, cum = set(), 0.0
        for k, (resolved, cost) in enumerate(per_rollout, 1):
            union |= resolved; cum += cost
            print(f"  {k:>2} {len(union):>6}/{n:<3} {100*len(union)/n:>7.0f}% {cum:>9.2f}")
        print("  (more rollouts → higher best-of-k resolution at a measured dollar cost)")

if __name__ == "__main__":
    main()