"""Tiny GPU-slot scheduler: runs a list of training configs across the local T4s,
N concurrent processes per GPU, skipping configs whose result file already exists."""

import argparse
import json
import os
import subprocess
import sys
import threading
import time
from queue import Queue

PY = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
                  ".venvr", "bin", "python")
TRAIN = os.path.join(os.path.dirname(os.path.abspath(__file__)), "train.py")


def worker(slot, gpu, q, logdir):
    while True:
        try:
            name, args = q.get_nowait()
        except Exception:
            return
        out = os.path.join(logdir, f"{name}.json")
        if os.path.exists(out) and os.path.getsize(out) > 0:
            print(f"[slot{slot}] skip {name}", flush=True)
            q.task_done()
            continue
        env = dict(os.environ, CUDA_VISIBLE_DEVICES=str(gpu), OMP_NUM_THREADS="4")
        cmd = [PY, TRAIN] + args + ["--out", out]
        log = os.path.join(logdir, f"{name}.log")
        t0 = time.time()
        print(f"[slot{slot}/gpu{gpu}] START {name}", flush=True)
        with open(log, "w") as fh:
            rc = subprocess.call(cmd, stdout=fh, stderr=subprocess.STDOUT, env=env)
        print(f"[slot{slot}/gpu{gpu}] DONE  {name} rc={rc} "
              f"({time.time()-t0:.0f}s)", flush=True)
        q.task_done()


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--jobs", required=True, help="json file: {name: [args...]}")
    ap.add_argument("--logdir", required=True)
    ap.add_argument("--gpus", default="0,1")
    ap.add_argument("--per-gpu", type=int, default=2)
    a = ap.parse_args()
    os.makedirs(a.logdir, exist_ok=True)
    jobs = json.load(open(a.jobs))
    q = Queue()
    for name, args in jobs.items():
        q.put((name, args))
    threads = []
    gpus = [g for g in a.gpus.split(",") if g]
    for gi, gpu in enumerate(gpus):
        for s in range(a.per_gpu):
            t = threading.Thread(target=worker, args=(f"{gi}.{s}", gpu, q, a.logdir))
            t.start()
            threads.append(t)
            time.sleep(3)
    for t in threads:
        t.join()
    print("ALL DONE", flush=True)


if __name__ == "__main__":
    main()