#!/usr/bin/env python3
"""Vast.ai client for the Igneum GPU fleet (6 October 2026). The key is read from ~/.config/vast/credentials and never
printed. Every rent and destroy is appended to ~/Desktop/fleet/ledger.jsonl with the offer's hourly price, so the spend
can be summed without the provider's statement.

  vast.py search "RTX 3090" [--n 8] [--min-cores 6] [--max-dph 0.5] [--gpus 1]
  vast.py rent <offer_id> --label <label> [--disk 60] [--image nvidia/cuda:12.8.1-devel-ubuntu24.04] [--onstart FILE]
  vast.py list                       every instance: id, label, gpu, state, ssh host and port, dph
  vast.py ssh <instance_id>          prints the ssh command line (host, port) for the fleet key
  vast.py destroy <instance_id>...   destroys and writes the ledger line
  vast.py spend                      the ledger's running total and the account's credit
"""
import json, os, sys, time, urllib.request, urllib.parse, urllib.error, argparse, datetime

KEY = open(os.path.expanduser("~/.config/vast/credentials")).read().strip()
BASE = "https://console.vast.ai/api"
LEDGER = os.path.expanduser("~/Desktop/fleet/ledger.jsonl")
IMAGE = "nvidia/cuda:12.8.1-devel-ubuntu24.04"

def call(method, path, body=None, timeout=90):
    req = urllib.request.Request(BASE + path, data=(json.dumps(body).encode() if body is not None else None),
                                 headers={"Authorization": "Bearer " + KEY, "Content-Type": "application/json", "Accept": "application/json"},
                                 method=method)
    for attempt in range(4):
        try:
            with urllib.request.urlopen(req, timeout=timeout) as r:
                raw = r.read() or b"{}"
            try: return json.loads(raw)
            except json.JSONDecodeError:
                if attempt < 3: time.sleep(3 * (attempt + 1)); continue
                raise SystemExit(f"{method} {path}: non-JSON reply: {raw[:120]!r}")
        except urllib.error.HTTPError as e:
            txt = e.read().decode(errors="replace")
            if e.code in (429, 502, 503, 504) and attempt < 3:
                time.sleep(3 * (attempt + 1)); continue
            raise SystemExit(f"{method} {path}: HTTP {e.code}: {txt[:400]}")
        except (urllib.error.URLError, TimeoutError) as e:
            if attempt < 3: time.sleep(3 * (attempt + 1)); continue
            raise

def now(): return datetime.datetime.now(datetime.timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")

def ledger(row):
    os.makedirs(os.path.dirname(LEDGER), exist_ok=True)
    with open(LEDGER, "a") as f: f.write(json.dumps(row) + "\n")

def search(gpu, n=8, min_cores=4, max_dph=None, gpus=1, min_disk=70, min_rel=0.95, min_down=200, cuda=12.8):
    q = {"gpu_name": {"eq": gpu}, "rentable": {"eq": True}, "verified": {"eq": True}, "num_gpus": {"eq": gpus},
         "cuda_max_good": {"gte": cuda}, "disk_space": {"gte": min_disk}, "reliability2": {"gte": min_rel},
         "inet_down": {"gte": min_down}, "cpu_cores_effective": {"gte": min_cores},
         "order": [["dph_total", "asc"]], "type": "on-demand", "limit": n}
    if max_dph: q["dph_total"] = {"lte": max_dph}
    d = call("GET", "/v0/bundles/?q=" + urllib.parse.quote(json.dumps(q)))
    return d.get("offers", [])

def fmt_offer(o):
    return (f"id {o['id']} ${o['dph_total']:.3f}/h {o.get('num_gpus')}x {o.get('gpu_name')} {o.get('gpu_ram')}MB cuda{o.get('cuda_max_good')} "
            f"cores{o.get('cpu_cores_effective',0):.0f} ram{o.get('cpu_ram',0)/1024:.0f}G disk{o.get('disk_space',0):.0f} dl{o.get('inet_down',0):.0f} "
            f"ul{o.get('inet_up',0):.0f} rel{o.get('reliability2',0):.3f} {o.get('geolocation')} drv{o.get('driver_version')}")

def rent(offer_id, label, disk=60, image=IMAGE, onstart=None, offer=None):
    body = {"client_id": "me", "image": image, "disk": disk, "label": label, "runtype": "ssh", "image_login": None,
            "python_utf8": False, "lang_utf8": False, "use_jupyter_lab": False, "cancel_unavail": True}
    if onstart: body["onstart"] = open(onstart).read()
    # the offer's price for the ledger
    o = offer or {}
    d = call("PUT", f"/v0/asks/{offer_id}/", body)
    if not d.get("success"): raise SystemExit(f"rent failed: {d}")
    iid = d.get("new_contract")
    row = {"t": now(), "event": "rent", "instance": iid, "offer": int(offer_id), "label": label, "gpu": o.get("gpu_name"),
           "num_gpus": o.get("num_gpus"), "dph": o.get("dph_total"), "disk": disk, "geo": o.get("geolocation")}
    ledger(row); print(json.dumps(row)); return iid

def instances():
    try:
        d = call("GET", "/v0/instances/?owner=me")
        return d.get("instances", [])
    except SystemExit:
        d = call("GET", "/v1/instances/")
        return d.get("instances", d if isinstance(d, list) else [])

def destroy(iid):
    try: d = call("DELETE", f"/v0/instances/{iid}/", {})
    except SystemExit: d = call("DELETE", f"/v1/instances/{iid}/", {})
    row = {"t": now(), "event": "destroy", "instance": int(iid), "ok": bool(d.get("success", True))}
    ledger(row); print(json.dumps(row))

def spend():
    rows = [json.loads(l) for l in open(LEDGER)] if os.path.exists(LEDGER) else []
    starts = {r["instance"]: r for r in rows if r["event"] == "rent"}
    stops = {r["instance"]: r for r in rows if r["event"] == "destroy"}
    total = 0.0; lines = []
    tnow = datetime.datetime.now(datetime.timezone.utc)
    for iid, r in starts.items():
        t0 = datetime.datetime.strptime(r["t"], "%Y-%m-%dT%H:%M:%SZ").replace(tzinfo=datetime.timezone.utc)
        t1 = datetime.datetime.strptime(stops[iid]["t"], "%Y-%m-%dT%H:%M:%SZ").replace(tzinfo=datetime.timezone.utc) if iid in stops else tnow
        h = max((t1 - t0).total_seconds(), 60) / 3600
        cost = h * float(r.get("dph") or 0) + (r.get("disk") or 0) * 0.0002 * h   # storage about USD 0.2 per TB-hour, approximate
        total += cost
        lines.append(f"{iid} {r['label']:<22} {r.get('gpu')!s:<16} {h:5.2f} h x ${r.get('dph') or 0:.3f} = ${cost:6.2f} {'running' if iid not in stops else 'destroyed'}")
    me = call("GET", "/v0/users/current/")
    print("\n".join(lines)); print(f"ledger total USD {total:.2f}; account credit USD {me.get('credit')}; balance {me.get('balance')}")

if __name__ == "__main__":
    ap = argparse.ArgumentParser(); sub = ap.add_subparsers(dest="cmd", required=True)
    s = sub.add_parser("search"); s.add_argument("gpu"); s.add_argument("--n", type=int, default=8); s.add_argument("--min-cores", type=int, default=4)
    s.add_argument("--max-dph", type=float); s.add_argument("--gpus", type=int, default=1); s.add_argument("--min-disk", type=int, default=70); s.add_argument("--json", action="store_true")
    r = sub.add_parser("rent"); r.add_argument("offer"); r.add_argument("--label", required=True); r.add_argument("--disk", type=int, default=60); r.add_argument("--image", default=IMAGE); r.add_argument("--onstart")
    sub.add_parser("list"); sub.add_parser("spend")
    d = sub.add_parser("destroy"); d.add_argument("ids", nargs="+")
    h = sub.add_parser("ssh"); h.add_argument("id")
    a = ap.parse_args()
    if a.cmd == "search":
        offs = search(a.gpu, a.n, a.min_cores, a.max_dph, a.gpus, a.min_disk)
        print(json.dumps(offs) if a.json else "\n".join(fmt_offer(o) for o in offs) or "no offers")
    elif a.cmd == "rent": rent(a.offer, a.label, a.disk, a.image, a.onstart)
    elif a.cmd == "list":
        for i in instances():
            print(f"{i.get('id')} {str(i.get('label')):<22} {i.get('num_gpus')}x {i.get('gpu_name')} {i.get('actual_status')}/{i.get('cur_state')} "
                  f"ssh {i.get('ssh_host')}:{i.get('ssh_port')} ${i.get('dph_total',0):.3f}/h {i.get('geolocation')} {i.get('status_msg') or ''}"[:200])
    elif a.cmd == "ssh":
        for i in instances():
            if str(i.get("id")) == a.id: print(f"ssh -i ~/.ssh/igneum-fleet -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null -p {i.get('ssh_port')} root@{i.get('ssh_host')}")
    elif a.cmd == "destroy":
        for i in a.ids: destroy(i)
    elif a.cmd == "spend": spend()