import json, os, time, subprocess
os.environ["HF_HUB_DISABLE_XET"] = "1"
from huggingface_hub import hf_hub_download

os.chdir("/home/corov/cyber/data")
tr = json.load(open("train_ids.json"))
ev = json.load(open("eval_ids.json"))

def files(rows):
    out = []
    for r in rows:
        tid = str(r["id"])
        for f in ("description.txt", "error.txt", "patch.diff"):
            out.append((f'data/{r["kind"]}/{tid}/{f}', r["kind"], tid, f))
    return out

todo = [("cgym_train", x) for x in files(tr)] + [("cgym_eval", x) for x in files(ev)]
failed = []
n_ok = 0
for i, (localdir, (fname, kind, tid, f)) in enumerate(todo):
    dest = os.path.join(localdir, "data", kind, tid, f)
    if os.path.exists(dest) and os.path.getsize(dest) > 0:
        n_ok += 1
        continue
    try:
        p = hf_hub_download("sunblaze-ucb/cybergym", fname, repo_type="dataset",
                            local_dir=localdir)
        n_ok += 1
    except Exception as e:
        failed.append(fname)
    if i % 60 == 59:
        time.sleep(4)

# one retry round for failures, slower
still = []
for fname in failed:
    localdir = "cgym_train" if any(fname.startswith(f"data/{r['kind']}/{str(r['id'])}/") for r in tr) else "cgym_eval"
    try:
        hf_hub_download("sunblaze-ucb/cybergym", fname, repo_type="dataset",
                        local_dir=localdir)
        n_ok += 1
        time.sleep(1)
    except Exception:
        still.append(fname)

print(f"OK={n_ok} FAILED={len(still)}")
for f in still[:10]:
    print("FAIL:", f)
for d in ("cgym_train", "cgym_eval"):
    n = subprocess.run(["find", d, "-name", "description.txt"],
                       capture_output=True, text=True).stdout.count("\n")
    print(d, "descriptions:", n)
print("FETCH2_DONE")