#!/usr/bin/env python3
"""The fleet orchestrator (Mac side). Registry: ~/Desktop/fleet/boxes.json ({instance: {...}}); per-box raw logs in
~/Desktop/fleet/<instance>/. Uses vast.py for the provider and ssh with ~/.ssh/igneum-fleet for the boxes.

  fleet.py rent-card "<gpu name>" <label> <archs> [--min-ram MB] [--max-ram MB] [--disk 60] [--phase 1]
  fleet.py wait [labels...]          until ssh answers on every (named) box
  fleet.py setup [labels...]         push the inputs and start box-setup.sh (nohup) on every (named) box
  fleet.py status [labels...]        the last RESULT or STAGE line of setup.log and the node's sync line
  fleet.py run <script> [labels...]  push tools/fleet/<script> and start it under nohup (box-matrix.sh, box-ember.sh, ...)
  fleet.py tail <label> [file]       the last 30 lines of a box's log
  fleet.py sh <label> "<cmd>"        run one command on a box
  fleet.py pull [labels...]          rsync /root/fleet/out and the logs into ~/Desktop/fleet/<instance>/
  fleet.py destroy <labels...>       pull, then destroy, then mark the registry
"""
import json, os, sys, subprocess, time, secrets, datetime
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import vast

HERE = os.path.dirname(os.path.abspath(__file__))
ROOT = os.path.expanduser("~/Desktop/fleet"); REG = os.path.join(ROOT, "boxes.json")
SSH_OPTS = ["-i", os.path.expanduser("~/.ssh/igneum-fleet"), "-o", "StrictHostKeyChecking=no", "-o", "UserKnownHostsFile=/dev/null",
            "-o", "LogLevel=ERROR", "-o", "ConnectTimeout=20", "-o", "ServerAliveInterval=15"]
INPUTS = [os.path.expanduser("~/Desktop/igneum-prove-wsl2-floor.zip"), os.path.join(HERE, "floor.patch"), os.path.join(HERE, "override.json")]

def now(): return datetime.datetime.now(datetime.timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")
import fcntl
def load():
    if not os.path.exists(REG): return {}
    for attempt in range(6):   # a reader can meet a half-written file only if a writer bypasses save(); retry anyway
        try: return json.load(open(REG))
        except json.JSONDecodeError:
            time.sleep(0.2 * (attempt + 1))
    raise
def save(reg):
    os.makedirs(ROOT, exist_ok=True); tmp = REG + f".tmp.{os.getpid()}"
    json.dump(reg, open(tmp, "w"), indent=1); os.replace(tmp, REG)
def patch(iid, **fields):
    with open(REG + ".lock", "w") as lk:
        fcntl.flock(lk, fcntl.LOCK_EX)
        reg = load(); reg.setdefault(str(iid), {}).update(fields); save(reg); return reg[str(iid)]

def boxes(labels, reg=None):
    reg = reg or load()
    sel = {k: v for k, v in reg.items() if v.get("state") != "destroyed" and (not labels or v["label"] in labels)}
    if labels and len(sel) != len(set(labels)): print("unknown or destroyed labels:", set(labels) - {v["label"] for v in sel.values()}, file=sys.stderr)
    return sel

def refresh_ssh(reg):
    live = {str(i.get("id")): i for i in vast.instances()}
    for iid, b in reg.items():
        if b.get("provider", "vast") == "vast" and iid in live:
            i = live[iid]; patch(iid, ssh_host=i.get("ssh_host"), ssh_port=i.get("ssh_port"), actual_status=i.get("actual_status"), status_msg=(i.get("status_msg") or "")[:120])
    if any(b.get("provider") == "runpod" and b.get("state") != "destroyed" for b in reg.values()):
        import runpod
        for p in runpod.pods():
            pid = str(p.get("id"))
            if pid in reg and reg[pid].get("state") != "destroyed":
                pm = p.get("portMappings") or {}
                patch(pid, ssh_host=p.get("publicIp") or None, ssh_port=pm.get("22") if p.get("publicIp") else None, actual_status=p.get("desiredStatus"), port_map=pm, status_msg="")
    return load()

def ssh(b, cmd, timeout=120, capture=True):
    try:
        r = subprocess.run(["ssh"] + SSH_OPTS + ["-p", str(b["ssh_port"]), f"root@{b['ssh_host']}", cmd], capture_output=capture, text=True, timeout=timeout, stdin=subprocess.DEVNULL)
    except subprocess.TimeoutExpired: return 124, "", "timeout"
    return r.returncode, (r.stdout or ""), (r.stderr or "")

def scp(b, files, dest):
    return subprocess.run(["scp"] + SSH_OPTS + ["-P", str(b["ssh_port"])] + files + [f"root@{b['ssh_host']}:{dest}"], capture_output=True, text=True, timeout=900).returncode

def rent_card(gpu, label, archs, min_ram=None, max_ram=None, disk=60, phase="1", min_cores=4, gpus=1):
    offers = vast.search(gpu, n=12, min_cores=min_cores, gpus=gpus)
    if min_ram: offers = [o for o in offers if (o.get("gpu_ram") or 0) >= min_ram]
    if max_ram: offers = [o for o in offers if (o.get("gpu_ram") or 0) <= max_ram]
    if not offers: print(f"{label}: no offer for {gpu}"); return None
    # prefer a few more cores for the build when the price is close: score = price + 0.01 per missing core under 8
    offers.sort(key=lambda o: o["dph_total"] + 0.01 * max(0, 8 - (o.get("cpu_cores_effective") or 0)))
    o = offers[0]; print(label, "->", vast.fmt_offer(o))
    iid = vast.rent(o["id"], label, disk, offer=o)
    reg = load()
    reg[str(iid)] = {"label": label, "card": gpu, "vram_mb": o.get("gpu_ram"), "archs": archs, "offer": o["id"], "dph": o["dph_total"], "num_gpus": o.get("num_gpus"),
                     "cores": o.get("cpu_cores_effective"), "ram_gb": round((o.get("cpu_ram") or 0) / 1024), "geo": o.get("geolocation"), "driver": o.get("driver_version"),
                     "provider": "vast", "phase": phase, "state": "renting", "rented_at": now(), "wallet": "0x" + secrets.token_hex(20)}
    save(reg); os.makedirs(os.path.join(ROOT, str(iid)), exist_ok=True); return iid

def wait(labels, limit=1200):
    t0 = time.time()
    while True:
        reg = refresh_ssh(load()); pending = []
        for iid, b in boxes(labels, reg).items():
            if b.get("state") in ("destroyed",): continue
            if b.get("ssh_ok"): continue
            if not b.get("ssh_host") or not b.get("ssh_port"): pending.append((b["label"], b.get("actual_status"), b.get("status_msg"))); continue
            try: rc, out, err = ssh(b, "nvidia-smi --query-gpu=name,memory.total --format=csv,noheader", timeout=40)
            except subprocess.TimeoutExpired: rc, out, err = 1, "", "timeout"
            if rc == 0 and out.strip(): patch(iid, ssh_ok=True, state="installing", gpu_seen=out.strip().replace("\n", ";")); print(f"{b['label']}: ssh ok, {out.strip()[:60]}")
            else: pending.append((b["label"], b.get("actual_status"), (err or out).strip()[:80]))
        if not pending: print("all up"); return
        if time.time() - t0 > limit: print("still pending:", pending); return
        print(f"waiting ({int(time.time()-t0)} s): " + "; ".join(f"{l} {s} {m}" for l, s, m in pending)); time.sleep(20)

def setup(labels):
    for iid, b in boxes(labels).items():
        if not b.get("ssh_ok"): print(b["label"], "no ssh yet"); continue
        ssh(b, "mkdir -p /root/fleet/in /root/fleet/out")
        if scp(b, INPUTS + [os.path.join(HERE, "box-setup.sh")], "/root/fleet/in/") != 0: print(b["label"], "scp failed"); continue
        rc, out, err = ssh(b, f"cd /root/fleet && chmod +x in/box-setup.sh && if pgrep -f '[b]ash in/box-setup.sh' >/dev/null; then echo already; else RUSTUP_TOOLCHAIN=stable ARCHS={b['archs']} LABEL={b['label']} WALLET={b['wallet']} setsid nohup in/box-setup.sh </dev/null >/dev/null 2>&1 & echo started; fi", timeout=60)
        print(b["label"], out.strip()[:40])
        patch(iid, setup_started=now())

def status(labels):
    reg = load()
    from concurrent.futures import ThreadPoolExecutor
    def one(item):
        iid, b = item
        if not b.get("ssh_ok"): return f"{b['label']:<14} {iid} no ssh", None
        rc, out, err = ssh(b, "grep -E '^(RESULT|STAGE)' /root/fleet/setup.log 2>/dev/null | tail -1; grep -c '^RESULT setup_done' /root/fleet/setup.log 2>/dev/null; /opt/igneum/pkg/bin/igneum-miner watch 1 grpc://127.0.0.1:26610 2>/dev/null | grep -o 'blocks=[0-9]*.*synced=[a-z]*' | tail -1; grep -E '^(RESULT|STAGE)' /root/fleet/out/matrix.log 2>/dev/null | tail -1; grep -E '^RESULT' /root/fleet/out/ember.log 2>/dev/null | tail -1; grep -E '^RESULT' /root/fleet/out/prover.log 2>/dev/null | tail -1", timeout=45)
        if rc == 124: return f"{b['label']:<14} {iid} ssh timeout", None
        lines = out.strip().split("\n")
        done = len(lines) > 1 and lines[1].strip() == "1"
        return f"{b['label']:<14} {iid} | " + " | ".join(l[:120] for l in lines if l.strip() and l.strip() not in ("0", "1")), done
    items = list(boxes(labels, reg).items())
    with ThreadPoolExecutor(max_workers=16) as ex: results = list(ex.map(one, items))
    for (iid, b), (line, done) in zip(items, results):
        print(line)
        if done and b.get("state") == "installing": patch(iid, state="running")

def run(script, labels, env=""):
    hub = next((v for v in load().values() if v.get("hub") and v.get("hub_peer")), None)
    if hub and "HUB_PEER" not in env: env = f"HUB_PEER={hub['hub_peer']} " + env
    for iid, b in boxes(labels).items():
        files = [os.path.join(HERE, script), os.path.join(HERE, "box-kill.sh")] + ([os.path.join(HERE, "box-prover.py")] if "prover" in script else [])
        if scp(b, files, "/root/fleet/in/") != 0: print(b["label"], "scp failed"); continue
        name = os.path.basename(script)
        logname = {"box-matrix.sh": "matrix", "box-ember.sh": "ember", "box-prover.sh": "prover"}.get(name, name)
        # rotate the stage's log before the start: an old run's end line must never be read as this run's (the 5070, 12:33Z)
        ssh(b, f"cd /root/fleet/out 2>/dev/null && for f in {logname}.log {logname}-launch.log rows.jsonl; do [ -f $f ] && mv $f $f.$(date +%s).old; done; true", timeout=30)
        ssh(b, f"cd /root/fleet && chmod +x in/{name} && LABEL={b['label']} WALLET={b['wallet']} ARCHS={b['archs']} {env} setsid nohup in/{name} </dev/null >/dev/null 2>&1 & echo ok", timeout=60)
        print(b["label"], name, "started")

def pull(labels):
    for iid, b in boxes(labels).items():
        d = os.path.join(ROOT, iid); os.makedirs(d, exist_ok=True)
        r = subprocess.run(["rsync", "-az", "-e", "ssh " + " ".join(SSH_OPTS) + f" -p {b['ssh_port']}", f"root@{b['ssh_host']}:/root/fleet/out/", f"root@{b['ssh_host']}:/root/fleet/setup.log", f"root@{b['ssh_host']}:/root/fleet/node.log", d + "/"], capture_output=True, text=True, timeout=600)
        print(b["label"], "pulled" if r.returncode == 0 else f"pull failed: {r.stderr[:200]}")

def destroy(labels):
    reg = load()
    for iid, b in boxes(labels, reg).items():
        if b.get("ssh_ok"):
            try: pull([b["label"]])
            except Exception as e: print("pull failed", e)
        vast.destroy(iid); t1 = now()
        t0 = datetime.datetime.strptime(b["rented_at"], "%Y-%m-%dT%H:%M:%SZ"); h = (datetime.datetime.strptime(t1, "%Y-%m-%dT%H:%M:%SZ") - t0).total_seconds() / 3600
        patch(iid, state="destroyed", destroyed_at=t1, hours=round(h, 2), cost_usd=round(h * b["dph"], 3)); print(b["label"], f"destroyed after {h:.2f} h, USD {h * b['dph']:.2f}")

if __name__ == "__main__":
    a = sys.argv[1:]
    if not a: print(__doc__); sys.exit(1)
    c = a[0]
    if c == "rent-card":
        kw = {}; pos = []
        i = 1
        while i < len(a):
            if a[i] == "--min-ram": kw["min_ram"] = int(a[i+1]); i += 2
            elif a[i] == "--max-ram": kw["max_ram"] = int(a[i+1]); i += 2
            elif a[i] == "--disk": kw["disk"] = int(a[i+1]); i += 2
            elif a[i] == "--phase": kw["phase"] = a[i+1]; i += 2
            elif a[i] == "--min-cores": kw["min_cores"] = int(a[i+1]); i += 2
            elif a[i] == "--gpus": kw["gpus"] = int(a[i+1]); i += 2
            else: pos.append(a[i]); i += 1
        rent_card(pos[0], pos[1], pos[2], **kw)
    elif c == "wait": wait(a[1:])
    elif c == "setup": setup(a[1:])
    elif c == "status": status(a[1:])
    elif c == "run": run(a[1], a[2:])
    elif c == "pull": pull(a[1:])
    elif c == "destroy": destroy(a[1:])
    elif c == "tail":
        b = [v for v in boxes([a[1]]).values()][0]; f = a[2] if len(a) > 2 else "/root/fleet/setup.log"; print(ssh(b, f"tail -n 30 {f}")[1])
    elif c == "sh":
        b = [v for v in boxes([a[1]]).values()][0]; rc, out, err = ssh(b, a[2], timeout=600); print(out, err)
    elif c == "dn2-version":   # the version string a Linux igneumd prints, read from its strings (the gate compares every box to it)
        import re as _re
        data = open(a[1], "rb").read(); m = _re.search(rb"igneumd/2\.1\.0-[0-9a-f]{8}", data); print(m.group(0).decode().replace("igneumd/", "igneumd_") if m else "")
    elif c == "list":
        for iid, b in load().items(): print(iid, b["label"], b["card"], b.get("state"), f"${b['dph']:.3f}/h", b.get("ssh_host"), b.get("ssh_port"))