"""Cell-queue job runner: runs every cell of an experiment script as a fresh subprocess (one solver thread each),
at most --workers concurrent, in phases (seed 0 for all cells first, then the remaining seeds) with cells ordered by
expected cost inside a phase, so that a --deadline cuts the number of seeds of the largest cells rather than whole
rows. Per-cell JSONs go to --cells-dir; cells already present are skipped (re-runnable). Finally --aggregate.

Usage: python cluster_driver.py --script preimage_sat.py --cells-dir results/cells_preimage \
           --out results/preimage_sat_scaling.json --workers 24 --deadline 7200 [--timeout 300]
"""
from __future__ import annotations
import argparse, json, os, pathlib, subprocess, sys, time
from concurrent.futures import ThreadPoolExecutor, as_completed

HERE = pathlib.Path(__file__).resolve().parent


def cell_name(c: dict) -> str:
    return "_".join(f"{k}-{c[k]}" for k in sorted(c))


def cost(c: dict) -> tuple:
    return (c.get("placement", "first") == "spread", c["r"], c.get("k", 0), c.get("iv") == "free", int(c.get("feedforward", 0)))


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--script", required=True)
    ap.add_argument("--cells-dir", required=True)
    ap.add_argument("--out", required=True)
    ap.add_argument("--workers", type=int, default=24)
    ap.add_argument("--deadline", type=float, default=1e9, help="seconds after start; no new cell is launched after it")
    ap.add_argument("--timeout", type=float, default=None)
    ap.add_argument("--list-args", default="")
    ap.add_argument("--python", default=sys.executable)
    ap.add_argument("--filter", default="", help="python expression over the cell dict c, e.g. \"c['k']==16 or c['r']<=16\"")
    a = ap.parse_args()
    script = str(HERE / a.script)
    cells_dir = pathlib.Path(a.cells_dir); cells_dir.mkdir(parents=True, exist_ok=True)
    lst = subprocess.run([a.python, script, "--list-cells"] + a.list_args.split(), capture_output=True, text=True, check=True)
    cells = [json.loads(l) for l in lst.stdout.splitlines() if l.strip()]
    if a.filter:
        cells = [c for c in cells if eval(a.filter, {}, {"c": c})]
    # phases: seed 0 first (all cells), then seeds 1, 2, ...; inside a phase by cost
    cells.sort(key=lambda c: (c["seed"], cost(c)))
    todo = [c for c in cells if not (cells_dir / (cell_name(c) + ".json")).exists()]
    print(f"{len(cells)} cells, {len(todo)} to run, workers={a.workers}, deadline={a.deadline}s", flush=True)
    t0 = time.time()
    log = open(cells_dir / "driver.log", "a")
    skipped = []

    def run(c):
        outp = cells_dir / (cell_name(c) + ".json")
        cmd = [a.python, script, "--cell", json.dumps(c), "--out", str(outp)]
        if a.timeout:
            cmd += ["--timeout", str(a.timeout)]
        env = dict(os.environ, OMP_NUM_THREADS="1", PYTHONPATH=str(HERE.parents[1] / "src"))
        t = time.time()
        try:
            p = subprocess.run(cmd, capture_output=True, text=True, env=env, timeout=(a.timeout or 600) + 120)
        except subprocess.TimeoutExpired:
            outp.write_text(json.dumps({"cell": c, "status": "timeout", "killed_by_driver": True,
                                        "wall_s": time.time() - t, "hostname": os.uname().nodename}))
            p = subprocess.CompletedProcess(cmd, -9, "", "killed by driver guard")
        status = "?"
        if outp.exists():
            try:
                status = json.loads(outp.read_text()).get("status")
            except Exception:
                status = "bad-json"
        else:
            status = "failed: " + p.stderr[-300:].replace("\n", " | ")
        line = f"{time.strftime('%H:%M:%S')} {cell_name(c)} -> {status} ({time.time()-t:.1f}s) [elapsed {time.time()-t0:.0f}s]"
        print(line, flush=True); log.write(line + "\n"); log.flush()
        return c, status

    with ThreadPoolExecutor(max_workers=a.workers) as ex:
        futs = []
        it = iter(todo)
        # submit lazily so the deadline can stop launches
        pending = set()
        for c in it:
            while len(pending) >= a.workers:
                done = next(as_completed(pending)); pending.discard(done)
            if time.time() - t0 > a.deadline:
                skipped.append(c)
                continue
            pending.add(ex.submit(run, c))
        for f in as_completed(pending):
            pass
    if skipped:
        (cells_dir / "skipped_cells.json").write_text(json.dumps(skipped, indent=1))
        print(f"deadline reached: {len(skipped)} cells not started (see skipped_cells.json)")
    subprocess.run([a.python, script, "--aggregate", str(cells_dir), "--out", a.out], check=True,
                   env=dict(os.environ, PYTHONPATH=str(HERE.parents[1] / "src")))
    print("done in %.0f s" % (time.time() - t0))


if __name__ == "__main__":
    main()