"""Minimal RunPod API v2 client for the H3 single-node benchmark pod."""

import argparse
import json
import os
import sys
import urllib.error
import urllib.parse
import urllib.request


API = "https://api.runpod.io/v2"
BLACKWELL_GPUS = (
    "NVIDIA RTX PRO 6000 Blackwell Server Edition",
    "NVIDIA RTX PRO 6000 Blackwell Workstation Edition",
    "NVIDIA RTX PRO 6000 Blackwell Max-Q Workstation Edition",
)
DEFAULT_IMAGE = "runpod/pytorch:1.1.0-cu1300-torch291-ubuntu2404"


def request(method: str, path: str, body=None, query=None):
    key = os.environ.get("RUNPOD_API_KEY")
    if not key:
        raise SystemExit("RUNPOD_API_KEY is required")
    url = f"{API}{path}"
    if query:
        url += "?" + urllib.parse.urlencode(query)
    data = None if body is None else json.dumps(body).encode("utf-8")
    call = urllib.request.Request(url, data=data, method=method)
    call.add_header("Authorization", f"Bearer {key}")
    call.add_header("Accept", "application/json")
    if data is not None:
        call.add_header("Content-Type", "application/json")
    try:
        with urllib.request.urlopen(call, timeout=60) as response:
            return json.load(response)
    except urllib.error.HTTPError as error:
        detail = error.read().decode("utf-8", errors="replace")
        raise SystemExit(f"RunPod API returned HTTP {error.code}: {detail}") from error


parser = argparse.ArgumentParser()
commands = parser.add_subparsers(dest="command", required=True)

catalog_parser = commands.add_parser("catalog")
catalog_parser.add_argument("--count", type=int, default=8)
catalog_parser.add_argument("--cloud", choices=("SECURE", "COMMUNITY"), default="SECURE")

create_parser = commands.add_parser("create")
create_parser.add_argument("--gpu", choices=BLACKWELL_GPUS, default=BLACKWELL_GPUS[0])
create_parser.add_argument("--count", type=int, default=8)
create_parser.add_argument("--cloud", choices=("SECURE", "COMMUNITY"), default="SECURE")
create_parser.add_argument("--datacenter")
create_parser.add_argument("--image", default=DEFAULT_IMAGE)
create_parser.add_argument("--disk", type=int, default=100)
create_parser.add_argument("--volume", type=int, default=100)
create_parser.add_argument("--yes", action="store_true")

get_parser = commands.add_parser("get")
get_parser.add_argument("pod_id")

terminate_parser = commands.add_parser("terminate")
terminate_parser.add_argument("pod_id")
terminate_parser.add_argument("--yes", action="store_true")

args = parser.parse_args()
if args.command == "catalog":
    response = request("GET", "/catalog/gpus", query={
        "include": "AVAILABILITY",
        "product": "POD",
        "count": args.count,
        "cloud": args.cloud,
        "minCudaVersion": "12.8",
    })
    response["gpus"] = [gpu for gpu in response["gpus"] if gpu["id"] in BLACKWELL_GPUS]
elif args.command == "create":
    if not args.yes:
        raise SystemExit("create rents billable GPUs; repeat with --yes after checking catalog")
    body = {
        "name": "h3-blackwell-distributed",
        "image": args.image,
        "gpu": {"id": args.gpu, "count": args.count, "minCudaVersion": "12.8"},
        "cloud": args.cloud,
        "disk": args.disk,
        "ports": ["22/tcp"],
        "mounts": {"persistent": {"size": args.volume, "path": "/workspace"}},
        "startSsh": True,
    }
    if args.datacenter:
        body["dataCenterIds"] = [args.datacenter]
    response = request("POST", "/pods", body=body)
elif args.command == "get":
    response = request("GET", f"/pods/{args.pod_id}")
else:
    if not args.yes:
        raise SystemExit("termination is irreversible; repeat with --yes")
    response = request("POST", f"/pods/{args.pod_id}/actions", body={"action": "terminate"})

json.dump(response, sys.stdout, indent=2)
sys.stdout.write("\n")