#!/usr/bin/env python3
"""Dispatch a job list over the MacLustr nodes (core-proportional shards) and collect results.
usage: dispatch_nodes.py launch --jobs results/manifests/jobs_main.json --name main [--nodes M3U96a:26,M3U96b:26,...]
       dispatch_nodes.py status --name main
       dispatch_nodes.py collect --name main        # rsync <node>:~/wp46-cache/forecasts_<name>/ -> results/forecasts/<name>/
Panels (daily returns + realized measures) must exist on every node under ~/wp46-cache/panels (synced here)."""
from __future__ import annotations

import argparse
import json
import subprocess
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
DEFAULT_NODES = "M3U96a:24,M3U96b:26,M2U64:22,M5M36:16,M4M64a:14,M4M64c:14"


def sh(cmd: list[str], timeout=600) -> str:
    return subprocess.run(cmd, capture_output=True, text=True, timeout=timeout).stdout


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("cmd", choices=["launch", "status", "collect", "sync-panels"])
    ap.add_argument("--jobs", default="")
    ap.add_argument("--name", required=True)
    ap.add_argument("--nodes", default=DEFAULT_NODES)
    a = ap.parse_args()
    nodes = [(n.split(":")[0], int(n.split(":")[1])) for n in a.nodes.split(",")]
    if a.cmd == "sync-panels":
        for n, _ in nodes:
            print(n, sh(["rsync", "-aq", "--delete", str(Path.home() / "wp46-cache" / "panels") + "/", f"{n}:~/wp46-cache/panels/"]))
        return
    if a.cmd == "launch":
        jobs = json.loads(Path(a.jobs).read_text())
        tot = sum(c for _, c in nodes)
        # proportional contiguous shards over the (heavy-first) job order -> interleave: node k takes jobs with
        # index % N in its share; simpler and balanced: assign job i to node by weighted round robin
        assign = {n: [] for n, _ in nodes}
        weights = [c / tot for _, c in nodes]
        acc = [0.0] * len(nodes)
        for i, j in enumerate(jobs):
            k = max(range(len(nodes)), key=lambda q: weights[q] * (i + 1) - acc[q])
            acc[k] += 1
            assign[nodes[k][0]].append(j)
        for n, cores in nodes:
            jf = ROOT / "results" / "manifests" / f"jobs_{a.name}_{n}.json"
            jf.write_text(json.dumps(assign[n]))
            subprocess.run(["scp", "-q", str(jf), f"{n}:~/wp46_uqo/results/manifests/"], check=True)
            print(n, cores, len(assign[n]), sh(["ssh", n, f"bash ~/wp46_uqo/scripts/node_run.sh results/manifests/jobs_{a.name}_{n}.json 0/1 {cores} {a.name}"]).strip())
        return
    if a.cmd == "status":
        for n, _ in nodes:
            out = sh(["ssh", n, f"cd ~/wp46_uqo && (wc -l < ~/wp46-cache/forecasts_{a.name}/_runs.jsonl 2>/dev/null || echo 0); "
                               f"python3 -c \"import json;print(len(json.load(open('results/manifests/jobs_{a.name}_{n}.json'))))\" 2>/dev/null; "
                               f"grep -c error ~/wp46-cache/forecasts_{a.name}/_runs.jsonl 2>/dev/null; uptime | sed 's/.*load/load/'; "
                               f"pgrep -fc run_benchmark.py"])
            print(n, " | ".join(x.strip() for x in out.splitlines()))
        return
    if a.cmd == "collect":
        dest = ROOT / "results" / "forecasts" / a.name
        dest.mkdir(parents=True, exist_ok=True)
        for n, _ in nodes:
            print(n, sh(["rsync", "-aq", f"{n}:~/wp46-cache/forecasts_{a.name}/", str(dest) + "/"], timeout=3600))
        runs = [json.loads(l) for f in sorted(dest.glob("_runs*.jsonl")) for l in f.read_text().splitlines()]
        print("collected runs:", len(runs), "errors:", sum("error" in r for r in runs))


if __name__ == "__main__":
    main()