"""Collect run logs into a small committable bundle.
On the training node after (some) runs finish:
    python scripts/collect.py --runs runs --out results/h200node1 && git add results && git commit -m "results" && git push
Checkpoints stay on the node (runs/ is gitignored); only JSONL summaries+curves are collected."""
import os, json, glob, argparse, shutil

p = argparse.ArgumentParser()
p.add_argument("--runs", default="runs")
p.add_argument("--out", default="results/local")
a = p.parse_args()
os.makedirs(a.out, exist_ok=True)
summary = {}
for d in sorted(glob.glob(os.path.join(a.runs, "*"))):
    name = os.path.basename(d)
    lp = os.path.join(d, "log_rank0.jsonl")
    if not os.path.exists(lp):
        continue
    rows = [json.loads(l) for l in open(lp)]
    meta = next((r for r in rows if r.get("kind") == "meta"), {})
    evals = [r for r in rows if r.get("kind") == "eval"]
    shutil.copy(lp, os.path.join(a.out, f"{name}.jsonl"))
    summary[name] = {"params": meta.get("params"), "world": meta.get("world"), "mode": (meta.get("cfg") or {}).get("mode"),
                     "n_probes": (meta.get("cfg") or {}).get("n_probes"), "done": os.path.exists(os.path.join(d, "DONE")),
                     "steps": evals[-1]["step"] if evals else 0, "val_loss": evals[-1]["val_loss"] if evals else None,
                     "tok_per_s": evals[-1].get("tok_per_s") if evals else None}
json.dump(summary, open(os.path.join(a.out, "summary.json"), "w"), indent=1)
print(json.dumps(summary, indent=1))