#!/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, patches) where patches maps each
    instance_id to the rollout's candidate patch text — needed so a no-leak
    consensus selector can vote across rollouts (see consensus_pick)."""
    preds, total_cost, patches = [], 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
        patches[iid] = patch
        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, patches

def _normalize_patch(patch):
    """Drop git blob-hash lines so identical *edits* compare equal regardless of
    incidental index metadata. Empty/whitespace patch → '' (an abstention)."""
    lines = [ln for ln in patch.splitlines() if not ln.startswith("index ")]
    return "\n".join(ln.rstrip() for ln in lines).strip()

def consensus_pick(patch_list):
    """No-leak BFT tournament vote over an instance's per-rollout candidate
    patches (ordered by rollout). Returns the 0-based rollout index of the
    plurality patch (ties → earliest rollout), or None if every rollout
    abstained (no patch). Never consults the hidden grader — pure agreement."""
    counts = {}
    for patch in patch_list:
        key = _normalize_patch(patch)
        if key:
            counts[key] = counts.get(key, 0) + 1
    if not counts:
        return None
    top = max(counts.values())
    for index, patch in enumerate(patch_list):
        key = _normalize_patch(patch)
        if key and counts[key] == top:
            return index
    return None

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, patches) 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:
        # Oracle (resolved-if-any) = the ceiling of any selection method.
        # Consensus = the no-leak BFT vote actually ships ONE patch per instance
        # (consensus_pick over the first k rollouts' candidates), so it reports
        # the resolution a *realizable* selector achieves. Same cost either way.
        print("  test-time-compute frontier (best-of-k):")
        print(f"  {'k':>2} {'oracle':>10} {'consensus':>11} {'cum $':>9}")
        cum = 0.0
        for k in range(1, args.rollouts + 1):
            cum += per_rollout[k - 1][1]
            union = set().union(*(per_rollout[i][0] for i in range(k)))
            selected = set()
            for iid in args.instances:
                candidates = [per_rollout[i][2].get(iid, "") for i in range(k)]
                pick = consensus_pick(candidates)
                if pick is not None and iid in per_rollout[pick][0]:
                    selected.add(iid)
            print(f"  {k:>2} {len(union):>6}/{n:<3} {len(selected):>7}/{n:<3} {cum:>9.2f}")
        print("  oracle = resolved-if-any (ceiling); consensus = the patch the BFT vote ships.")
        print("  gap = what better selection could still capture; cost is identical.")

if __name__ == "__main__":
    main()