#!/usr/bin/env python
"""Run a plan of generation jobs on a pool of GPUs, one vLLM worker per GPU.

    python scripts/run_jobs.py plans/trainer.json --gpus 0-7

A plan lists runs: {"name", "model", "data", "n", "shards", "est" (millions of tokens, for ordering),
optional "priority" (higher starts first), "preset", "set" and "args"}. "model" is "base", a folder name under
models/ ("qwen35-p50c") or a path; "data" is a file name under data/ without .jsonl. Each run is split into
shards; the biggest shards start first. Finished shards (a .done.json) are skipped, so the script can be restarted
any time. When every shard of a run is done, the run is merged into runs/<name>/merged.jsonl and summarized in
summary.json.
"""
import argparse
import fnmatch
import json
import os
import subprocess
import sys
import time

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, ROOT)
from wsp.summary import summarize  # noqa: E402

ap = argparse.ArgumentParser()
ap.add_argument("plan")
ap.add_argument("--gpus", default="0-7")
ap.add_argument("--lab", default="/nvme/wsp-lab")
ap.add_argument("--only", nargs="*", help="run names or globs")
ap.add_argument("--retries", type=int, default=1)
ap.add_argument("--dry", action="store_true")
a = ap.parse_args()

if "-" in a.gpus:
    lo, hi = map(int, a.gpus.split("-"))
    GPUS = list(range(lo, hi + 1))
else:
    GPUS = [int(g) for g in a.gpus.split(",")]

plan = json.load(open(a.plan))
runs = [r for r in plan["runs"] if not a.only or any(fnmatch.fnmatch(r["name"], p) for p in a.only)]


def model_path(m):
    if m == "base":
        return os.path.join(a.lab, "models", "Qwen3.5-0.8B")
    return m if m.startswith("/") else os.path.join(a.lab, "models", m)


def run_dir(r):
    return os.path.join(a.lab, "runs", r["name"])


def shard_out(r, i):
    return os.path.join(run_dir(r), f"shard{i}of{r['shards']}")


def merge(r):
    rows = [json.loads(line) for line in open(os.path.join(a.lab, "data", r["data"] + ".jsonl"))]
    recs, done = {}, []
    for i in range(r["shards"]):
        for line in open(shard_out(r, i) + ".jsonl"):
            x = json.loads(line)
            recs[x["id"]] = x
        done.append(json.load(open(shard_out(r, i) + ".done.json")))
    missing = [x["id"] for x in rows if x["id"] not in recs]
    if missing and "--limit" not in r.get("args", []):  # with --limit, only the first rows of each shard ran
        raise KeyError(f"{len(missing)} problems have no answers, the first is {missing[0]}")
    merged = [recs[x["id"]] for x in rows if x["id"] in recs]
    with open(os.path.join(run_dir(r), "merged.jsonl"), "w") as f:
        for x in merged:
            f.write(json.dumps(x) + "\n")
    s = summarize(merged)
    s["run"] = r
    s["sampling"] = done[0]["sampling"]
    s["edit_note"] = done[0]["edit_note"]
    s["tokens_total"] = sum(d["tokens"] for d in done)
    s["gpu_seconds"] = round(sum(d["gen_s"] + d["load_s"] + d["grade_s"] for d in done))
    s["tok_per_s_per_gpu"] = round(sum(d["tokens"] for d in done) / max(sum(d["gen_s"] for d in done), 1e-9))
    json.dump(s, open(os.path.join(run_dir(r), "summary.json"), "w"), indent=1)
    print(f"[{time.strftime('%H:%M:%S')}] MERGED {r['name']}: acc={s['acc']:.3f} cut={s['cut_share']:.3f} "
          f"mean_tokens={s['mean_tokens']:.0f} tok/s/gpu={s['tok_per_s_per_gpu']}", flush=True)


units = []
for r in runs:
    for i in range(r["shards"]):
        if not os.path.exists(shard_out(r, i) + ".done.json"):
            units.append((r, i))
units.sort(key=lambda u: (-u[0].get("priority", 0), -u[0].get("est", 1) / u[0]["shards"]))
print(f"{len(units)} shards to run on GPUs {GPUS}:", flush=True)
for r, i in units:
    print(f"  {r['name']} {i}/{r['shards']} ~{r.get('est', 0) / r['shards']:.1f}M tokens", flush=True)
if a.dry:
    sys.exit(0)

L = a.lab
env = dict(os.environ, NLTK_DATA=f"{L}/nltk_data", HF_HOME=f"{L}/hf", VLLM_CACHE_ROOT=f"{L}/vllm-cache",
           FLASHINFER_WORKSPACE_BASE=L, TRITON_CACHE_DIR=f"{L}/triton-cache",
           TORCHINDUCTOR_CACHE_DIR=f"{L}/inductor-cache", XDG_CACHE_HOME=f"{L}/xdg-cache",
           MPLCONFIGDIR=f"{L}/mpl", TOKENIZERS_PARALLELISM="false", PYTHONUNBUFFERED="1")
free, running, tries = list(GPUS), {}, {}
failed = []


def launch(r, i, g):
    os.makedirs(run_dir(r), exist_ok=True)
    cmd = [sys.executable, os.path.join(ROOT, "scripts", "gen_worker.py"), "--model", model_path(r["model"]),
           "--data", os.path.join(a.lab, "data", r["data"] + ".jsonl"), "--out", shard_out(r, i),
           "--n", str(r["n"]), "--shard", f"{i}/{r['shards']}", "--preset", r.get("preset", "think35")]
    for kv in r.get("set", []):
        cmd += ["--set", kv]
    cmd += r.get("args", [])
    log = open(shard_out(r, i) + ".log", "a")
    p = subprocess.Popen(cmd, env=dict(env, CUDA_VISIBLE_DEVICES=str(g)), stdout=log, stderr=subprocess.STDOUT)
    running[p] = (r, i, g, time.time())
    print(f"[{time.strftime('%H:%M:%S')}] START {r['name']} {i}/{r['shards']} on GPU {g}", flush=True)


while units or running:
    while units and free:
        r, i = units.pop(0)
        launch(r, i, free.pop(0))
    time.sleep(10)
    for p in list(running):
        if p.poll() is None:
            continue
        r, i, g, t0 = running.pop(p)
        free.append(g)
        ok = p.returncode == 0 and os.path.exists(shard_out(r, i) + ".done.json")
        mins = (time.time() - t0) / 60
        print(f"[{time.strftime('%H:%M:%S')}] {'DONE' if ok else 'FAIL'} {r['name']} {i}/{r['shards']} "
              f"on GPU {g} after {mins:.1f} min (exit {p.returncode})", flush=True)
        if not ok:
            k = (r["name"], i)
            tries[k] = tries.get(k, 0) + 1
            if tries[k] <= a.retries:
                units.insert(0, (r, i))
            else:
                failed.append(k)
            continue
        if all(os.path.exists(shard_out(r, j) + ".done.json") for j in range(r["shards"])):
            try:
                merge(r)
            except Exception as e:  # noqa: BLE001
                print("MERGE FAILED", r["name"], repr(e), flush=True)

for r in runs:
    if all(os.path.exists(shard_out(r, j) + ".done.json") for j in range(r["shards"])) and \
            not os.path.exists(os.path.join(run_dir(r), "summary.json")):
        merge(r)
print("ALL DONE" if not failed else f"FINISHED WITH FAILURES: {failed}", flush=True)