"""Run a prepared jobs file through the Rust engine (optionally K engine processes in parallel, each job single-threaded)
and record provenance on the executing host. Stdlib only, Python ≥ 3.9.

Usage: python3 run_batch_remote.py --engine ./rho_engine --jobs study.jobs.json --out study_<host>.json [--parallel K] [--params params.json]
"""
import argparse, json, pathlib, subprocess, sys, time, os
sys.path.insert(0, str(pathlib.Path(__file__).resolve().parents[2] / "src"))
from wp34.provenance import record


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--engine", required=True); ap.add_argument("--jobs", required=True); ap.add_argument("--out", required=True)
    ap.add_argument("--params", default=None); ap.add_argument("--parallel", type=int, default=1); ap.add_argument("--seed", type=int, default=34)
    a = ap.parse_args(); t0 = time.time()
    jobs = json.loads(pathlib.Path(a.jobs).read_text())
    out = pathlib.Path(a.out); out.parent.mkdir(parents=True, exist_ok=True)
    K = max(1, a.parallel)
    chunks = [jobs[i::K] for i in range(K)]
    procs, files = [], []
    for i, ch in enumerate(chunks):
        if not ch:
            continue
        jf = out.with_suffix(f".chunk{i}.jobs.json"); rf = out.with_suffix(f".chunk{i}.raw.json")
        jf.write_text(json.dumps(ch)); files.append((i, rf))
        procs.append(subprocess.Popen([a.engine, "batch", "--file", str(jf), "--out", str(rf)], stdout=subprocess.DEVNULL, stderr=open(out.with_suffix(f".chunk{i}.err"), "w")))
    for p in procs:
        p.wait()
    # merge in original order
    merged = [None] * len(jobs)
    for i, rf in files:
        res = json.loads(rf.read_text())
        for k, r in enumerate(res):
            merged[i + k * K] = r
    ok = sum(1 for r in merged if r and r.get("ok"))
    params = json.loads(pathlib.Path(a.params).read_text()) if a.params else {}
    params.update({"engine": a.engine, "jobs_file": a.jobs, "parallel_processes": K})
    record("ecdlp_study", params, {"n_jobs": len(jobs), "n_ok": ok, "runs": merged}, seed=a.seed, out=str(out), started=t0)
    for i, rf in files:
        rf.unlink(); out.with_suffix(f".chunk{i}.jobs.json").unlink()
    print(f"done: {ok}/{len(jobs)} verified in {time.time()-t0:.0f}s on {os.uname().nodename} → {out}")


if __name__ == "__main__":
    main()