"""Recree le pod Ornith A40x2 depuis l'image precompilee du Space HF.

Pourquoi recreer et non redemarrer : Runpod a refuse le `start` avec
"not enough free GPUs on the host machine" -- un pod arrete est lie a son hote,
et l'hote n'a plus deux A40 libres. Le recreer le replace sur un hote qui en a.

Pourquoi un template : `create-pod` n'a pas de champ pour la commande de
demarrage ; elle passe par `dockerStartCmd` du template. Et l'image vient du
registre HF (prive), d'ou `containerRegistryAuthId`.
"""
import json
import os
import sys
import urllib.error
import urllib.request

ICI = os.path.dirname(os.path.abspath(__file__))
K = open(os.path.join(ICI, "rpk")).read().strip()
AUTH = open(os.path.join(ICI, "regauth")).read().strip()
HF_TOKEN = open(os.path.expanduser("~/.cache/huggingface/token")).read().strip()
PUB = open(os.path.join(ICI, "ssh", "pod_key.pub")).read().strip()
PUB_USER = ("ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAIKK2mVCNr6s2S+r9fpX0/SP1gZiceL0pg/"
            "NQCTTBA8Rs eira.patrick@gmail.com")
IMAGE = os.environ.get("POD_IMAGE", "registry.hf.space/patdev-ornith-vllm-a40:latest")
GPU = int(os.environ.get("POD_GPUS", "2"))


def rp(path, method="GET", body=None):
    r = urllib.request.Request(
        "https://rest.runpod.io/v1" + path, method=method,
        headers={"Authorization": "Bearer " + K, "Content-Type": "application/json"},
        data=json.dumps(body).encode() if body else None)
    try:
        return json.loads(urllib.request.urlopen(r, timeout=120).read().decode())
    except urllib.error.HTTPError as e:
        return {"err": e.code, "body": e.read().decode()[:500]}


ENV = {
    "HF_TOKEN": HF_TOKEN,
    "HF_XET_HIGH_PERFORMANCE": "1",
    "HF_HUB_ENABLE_HF_TRANSFER": "0",
    "VL_MODEL": "ornith",
    "PUBLIC_KEY": PUB_USER + "\n" + PUB,
    # Les serveurs MCP node meurent avec un preload herite ; on le neutralise.
    "NODE_OPTIONS": "",
}

tpl_name = "ornith-vllm-a40-prebaked"
tpls = rp("/templates")
tpl = next((t for t in tpls if t.get("name") == tpl_name), None) if isinstance(tpls, list) else None
if tpl and tpl.get("imageName") != IMAGE:
    rp(f"/templates/{tpl['id']}", "DELETE")
    tpl = None
if not tpl:
    tpl = rp("/templates", "POST", {
        "name": tpl_name,
        "imageName": IMAGE,
        "containerRegistryAuthId": AUTH,
        "containerDiskInGb": 200,
        "volumeInGb": 0,
        "ports": ["8080/http", "22/tcp"],
        "env": ENV,
        # /start_pod.sh = sshd en arriere-plan + bootstrap du Hub (rechargement a chaud)
        "dockerStartCmd": ["/start_pod.sh"],
        "isPublic": False,
    })
    if "err" in tpl:
        sys.exit(f"template: {tpl}")
print("template", tpl["id"], tpl.get("imageName"))

# CUDA 13.0 epingle : sans lui, le meme deploiement marche ou casse selon le
# tirage d'hote. cuda-compat couvre 12.x, mais autant ne pas jouer.
corps = {
    "name": "ornith-a40x2",
    "templateId": tpl["id"],
    "gpuTypeIds": ["NVIDIA A40"],
    "gpuCount": GPU,
    "cloudType": "SECURE",
    "containerDiskInGb": 200,
    "volumeInGb": 0,
    "ports": ["8080/http", "22/tcp"],
    "env": ENV,
    "allowedCudaVersions": ["13.0", "12.8"],
    "supportPublicIp": True,
}
essais = [dict(corps), {**corps, "dataCenterIds": ["EU-SE-1"]},
          {**corps, "dataCenterIds": ["EU-RO-1"]}, {**corps, "dataCenterIds": ["EU-CZ-1"]},
          {**corps, "cloudType": "COMMUNITY"}]
for i, c in enumerate(essais, 1):
    p = rp("/pods", "POST", c)
    if "err" not in p:
        print(json.dumps({x: p.get(x) for x in ("id", "name", "desiredStatus", "costPerHr",
                                               "gpuCount", "cudaVersion", "machineId")}, indent=1))
        open(os.path.join(ICI, "podid"), "w").write(p["id"])
        print("PODID", p["id"])
        break
    print(f"essai {i} refuse : {p['body'][:200]}")
else:
    sys.exit("aucun placement possible")