"""Fetch instance metadata and each repository's file list at the instance's base commit.

Sources:
- SWE-bench Lite test split (princeton-nlp/SWE-bench_Lite): repo, base_commit, patch.
- Multi-SWE-bench Flash (ByteDance-Seed/Multi-SWE-bench-flash, split train): org/repo,
  base.sha, fix_patch, language.

Gold files are the files touched by the resolving patch (``diff --git`` headers), which
is the harness's definition for both benchmarks.

    uv run --no-project --with pyarrow python fetch_trees.py

Writes raw/instances.json ({instance_id: {benchmark, repo, base_commit, gold, language}})
and raw/trees/<instance_id>.json ({truncated, paths}) through the authenticated `gh` CLI
(one git-trees API call per instance, 5000 requests/hour).
"""
import json, pathlib, subprocess, sys, urllib.request

RAW = pathlib.Path(__file__).parent / "raw"
LITE_PARQUET = "https://huggingface.co/datasets/princeton-nlp/SWE-bench_Lite/resolve/main/data/test-00000-of-00001.parquet"
FLASH_JSONL = "https://huggingface.co/datasets/ByteDance-Seed/Multi-SWE-bench-flash/resolve/main/multi_swe_bench_flash.jsonl"
LANG = {"c": "C", "c++": "C++", "go": "Go", "java": "Java", "javascript": "JavaScript",
        "rust": "Rust", "typescript": "TypeScript", "python": "Python"}


def patch_files(patch: str) -> list[str]:
    return sorted({l.split(" b/")[-1].strip() for l in patch.splitlines() if l.startswith("diff --git")})


def load_instances() -> dict:
    out = {}
    pq_path = RAW / "lite_test.parquet"
    if not pq_path.exists():
        urllib.request.urlretrieve(LITE_PARQUET, pq_path)
    import pyarrow.parquet as pq
    for r in pq.read_table(pq_path, columns=["instance_id", "repo", "base_commit", "patch"]).to_pylist():
        out[r["instance_id"]] = {"benchmark": "SWE-bench Lite", "repo": r["repo"],
                                 "base_commit": r["base_commit"], "gold": patch_files(r["patch"]),
                                 "language": "Python"}
    fl_path = RAW / "multi_swe_bench_flash.jsonl"
    if not fl_path.exists():
        urllib.request.urlretrieve(FLASH_JSONL, fl_path)
    for line in open(fl_path):
        r = json.loads(line)
        out[r["instance_id"]] = {"benchmark": "Multi-SWE-bench Flash", "repo": f"{r['org']}/{r['repo']}",
                                 "base_commit": r["base"]["sha"], "gold": patch_files(r["fix_patch"]),
                                 "language": LANG.get(r["language"], r["language"])}
    return out


def main():
    (RAW / "trees").mkdir(parents=True, exist_ok=True)
    inst = load_instances()
    json.dump(inst, open(RAW / "instances.json", "w"), indent=0)
    only = set(sys.argv[1:])
    for iid, r in sorted(inst.items()):
        if only and iid not in only:
            continue
        out = RAW / "trees" / f"{iid}.json"
        if out.exists():
            continue
        res = subprocess.run(
            ["gh", "api", f"repos/{r['repo']}/git/trees/{r['base_commit']}?recursive=1", "--jq",
             '{truncated: .truncated, paths: [.tree[] | select(.type=="blob") | .path]}'],
            capture_output=True, text=True, timeout=120)
        if res.returncode != 0:
            print("FAIL", iid, res.stderr[:200], flush=True); continue
        out.write_text(res.stdout)
        if json.loads(res.stdout)["truncated"]:
            print("TRUNCATED", iid, flush=True)
    print("done", len(inst))


if __name__ == "__main__":
    main()