#!/usr/bin/env python3
"""Push upstream API keys into SwarmUI (user GenericData via SetAPIKey).

Stdlib only. Keys file: /tmp/gpu-rent-swarm-api-keys.json (mode 600), shape:
  {"civitai_api": "...", "huggingface_api": "..."}  # omit empty
"""

from __future__ import annotations

import json
import sys
import time
import urllib.error
import urllib.request
from pathlib import Path

KEYS_PATH = Path("/tmp/gpu-rent-swarm-api-keys.json")
SWARM = "http://127.0.0.1:7801"
ACCEPTED = ("civitai_api", "huggingface_api", "stability_api")


def post(path: str, payload: dict, timeout: float = 15.0) -> dict:
    body = json.dumps(payload).encode("utf-8")
    req = urllib.request.Request(
        f"{SWARM}{path}",
        data=body,
        headers={"Content-Type": "application/json"},
        method="POST",
    )
    with urllib.request.urlopen(req, timeout=timeout) as resp:
        return json.loads(resp.read().decode("utf-8"))


def wait_session(deadline: float) -> str:
    last = ""
    while time.time() < deadline:
        try:
            data = post("/API/GetNewSession", {})
            sid = str(data.get("session_id") or "")
            if sid:
                return sid
            last = "no session_id"
        except (urllib.error.URLError, urllib.error.HTTPError, TimeoutError, OSError, json.JSONDecodeError) as exc:
            last = str(exc)[:160]
        time.sleep(2)
    raise SystemExit(f"SwarmUI session unavailable: {last}")


def main() -> int:
    if not KEYS_PATH.is_file():
        print("no keys file — skip")
        return 0
    try:
        raw = json.loads(KEYS_PATH.read_text(encoding="utf-8"))
    except (OSError, json.JSONDecodeError) as exc:
        print(f"bad keys file: {exc}")
        return 1
    finally:
        try:
            KEYS_PATH.unlink(missing_ok=True)
        except OSError:
            pass

    keys = {k: str(v).strip() for k, v in (raw or {}).items() if k in ACCEPTED and str(v).strip()}
    if not keys:
        print("no api keys to set")
        return 0

    sid = wait_session(time.time() + 120)
    for key_type, value in keys.items():
        try:
            resp = post(
                "/API/SetAPIKey",
                {"session_id": sid, "keyType": key_type, "key": value},
            )
        except (urllib.error.URLError, urllib.error.HTTPError, TimeoutError, OSError, json.JSONDecodeError) as exc:
            print(f"SetAPIKey {key_type} failed: {exc}")
            return 1
        if resp.get("error"):
            print(f"SetAPIKey {key_type}: {resp['error']}")
            return 1
        if not resp.get("success"):
            print(f"SetAPIKey {key_type}: unexpected {resp}")
            return 1
        print(f"SetAPIKey {key_type}=ok")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())