#!/usr/bin/env python3
"""RunPod REST client for the fleet (https://rest.runpod.io/v1, Bearer from ~/.config/runpod/credentials; the GraphQL
endpoint refuses the Bearer header, so nothing here touches account settings: the fleet key goes in per pod through
the PUBLIC_KEY env, which RunPod's images write to authorized_keys). Ledger lines as vast.py.

  runpod.py gpus                      the GPU types with availability and prices
  runpod.py rent <gpuTypeId> <label> [--count 8] [--disk 80] [--image ...] [--secure]
  runpod.py list                      pods: id, name, status, ssh endpoint
  runpod.py destroy <pod_id>...
"""
import json, os, sys, time, urllib.request, urllib.error, datetime, argparse
KEY = open(os.path.expanduser("~/.config/runpod/credentials")).read().strip()
BASE = "https://rest.runpod.io/v1"
LEDGER = os.path.expanduser("~/Desktop/fleet/ledger.jsonl")
IMAGE = "runpod/pytorch:2.8.0-py3.11-cuda12.8.1-cudnn-devel-ubuntu22.04"
PUB = open(os.path.expanduser("~/.ssh/igneum-fleet.pub")).read().strip()
def call(method, path, body=None):
    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", "User-Agent": "curl/8.7.1", "Accept": "*/*"}, method=method)   # Cloudflare 1010 refuses urllib's default agent
    try:
        with urllib.request.urlopen(req, timeout=90) as r:
            t = r.read(); return json.loads(t) if t else {}
    except urllib.error.HTTPError as e: raise SystemExit(f"{method} {path}: HTTP {e.code}: {e.read().decode(errors='replace')[:400]}")
def now(): return datetime.datetime.now(datetime.timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")
def ledger(row):
    with open(LEDGER, "a") as f: f.write(json.dumps(row) + "\n")
def gpus():
    for g in call("GET", "/gputypes"):
        print(f"{g.get('id'):<34} {g.get('displayName','')!s:<22} {g.get('memoryInGb')}GB secure ${g.get('securePrice')}/h community ${g.get('communityPrice')}/h max {g.get('maxGpuCount')} avail {g.get('lowestPrice', {}) if isinstance(g.get('lowestPrice'), dict) else ''}")
def rent(gpu, label, count=8, disk=80, image=IMAGE, secure=False, ports=("22/tcp",)):
    body = {"name": label, "imageName": image, "gpuTypeIds": [gpu], "gpuCount": count, "containerDiskInGb": disk, "volumeInGb": 0,
            "cloudType": "SECURE" if secure else "COMMUNITY", "ports": list(ports), "env": {"PUBLIC_KEY": PUB}, "supportPublicIp": True, "computeType": "GPU"}
    d = call("POST", "/pods", body)
    row = {"t": now(), "event": "rent", "provider": "runpod", "instance": d.get("id"), "label": label, "gpu": gpu, "num_gpus": count, "dph": d.get("costPerHr"), "disk": disk}
    ledger(row); print(json.dumps(row)); print(json.dumps(d)[:600]); return d
def pods():
    out = call("GET", "/pods")
    return out if isinstance(out, list) else out.get("pods", [])
def ssh_of(p):
    for m in (p.get("portMappings") or {}).items() if isinstance(p.get("portMappings"), dict) else []:
        if str(m[0]) == "22": return p.get("publicIp"), m[1]
    return p.get("publicIp"), None
def destroy(pid):
    d = call("DELETE", f"/pods/{pid}"); ledger({"t": now(), "event": "destroy", "provider": "runpod", "instance": pid}); print("destroyed", pid, d)
if __name__ == "__main__":
    a = sys.argv[1:]
    if a[0] == "gpus": gpus()
    elif a[0] == "rent":
        ap = argparse.ArgumentParser(); ap.add_argument("gpu"); ap.add_argument("label"); ap.add_argument("--count", type=int, default=8); ap.add_argument("--disk", type=int, default=80); ap.add_argument("--image", default=IMAGE); ap.add_argument("--secure", action="store_true")
        n = ap.parse_args(a[1:]); rent(n.gpu, n.label, n.count, n.disk, n.image, n.secure)
    elif a[0] == "list":
        for p in pods(): print(p.get("id"), p.get("name"), p.get("desiredStatus"), p.get("gpuCount"), p.get("machine", {}).get("gpuTypeId") if isinstance(p.get("machine"), dict) else "", f"${p.get('costPerHr')}/h", ssh_of(p), (p.get("portMappings")))
    elif a[0] == "destroy":
        for i in a[1:]: destroy(i)
    elif a[0] == "raw": print(json.dumps(pods(), indent=1)[:3000])