"""Create (or reuse) the RunPod serverless endpoint for the docling parse worker.

The endpoint configuration lives here as code so it is reproducible: the
5090-and-up GPU tiers, hosts restricted to CUDA 13.0 drivers (matching the
image's torch build), zero active workers, FlashBoot off, and the GHCR pull
credential stored once under a named registry auth.

Idempotent by name: each piece (registry auth, template, endpoint) is
reused when it already exists. Templates are versioned by image tag, so
pointing the endpoint at a new image re-runs with the new tag and moves the
endpoint to the new template.

Usage:
    uv run python runpod/create_endpoint.py --image ghcr.io/xmandeng/quber-runpod:v0.1.0

Environment:
    RUNPOD_API_KEY   required
    GHCR_USERNAME    required only when the registry auth is first created
    GHCR_PULL_TOKEN  required only when the registry auth is first created
                     (a long-lived PAT with read:packages)
"""

from __future__ import annotations

import os
import sys
import time
from typing import Any

import click
import httpx

API_BASE = "https://rest.runpod.io/v1"
GRAPHQL_URL = "https://api.runpod.io/graphql"

# 5090 first, with 48GB datacenter Ampere cards as fallbacks. This is the
# selection proven live: the consumer 24GB pools repeatedly placed workers
# on hosts whose drivers were too old for the image's torch build.
GPU_TYPE_IDS = [
    "NVIDIA GeForce RTX 5090",
    "NVIDIA A40",
    "NVIDIA RTX A6000",
]

# The image ships torch built for CUDA 13.0, which refuses to initialize on
# hosts whose driver only supports an older CUDA ("driver too old", and
# torch then reports no usable device). RunPod's scheduler default accepts
# hosts down to CUDA 11.8, so without this filter jobs land on 12.4/12.8
# hosts and every worker fails its GPU fitness check. Keep in lockstep with
# the torch CUDA build in uv.lock.
ALLOWED_CUDA_VERSIONS = ["13.0"]

REGISTRY_AUTH_NAME = "ghcr-quber"
ENDPOINT_NAME = "quber-parse"

# RunPod injects RTX PRO 6000 Blackwell MIG slices into the serverless
# capacity categories by default (1g.24gb into 24GB, 2g.48gb into 48GB) and
# the scheduler works on categories, so pinning gpuTypeIds does not keep
# them out. The parse pipeline fails on those slices. The REST API has no
# field for MIG at all; the exclusion lives in the GraphQL gpuIds string as
# a negative entry — the same representation the console's opt-out checkbox
# writes.
MIG_EXCLUSIONS = [
    "NVIDIA RTX PRO 6000 Blackwell Server Edition MIG 1g.24gb",
    "NVIDIA RTX PRO 6000 Blackwell Server Edition MIG 2g.48gb",
]
WORKER_DRAIN_TIMEOUT_SECONDS = 180


def api_client() -> httpx.Client:
    api_key = os.environ.get("RUNPOD_API_KEY")
    if not api_key:
        click.echo("RUNPOD_API_KEY is not set", err=True)
        sys.exit(1)
    return httpx.Client(
        base_url=API_BASE,
        headers={"Authorization": f"Bearer {api_key}"},
        timeout=60,
    )


def ensure_registry_auth(http: httpx.Client) -> str:
    response = http.get("/containerregistryauth")
    response.raise_for_status()
    for auth in response.json() or []:
        if auth["name"] == REGISTRY_AUTH_NAME:
            click.echo(f"Registry auth exists: {auth['id']}")
            return auth["id"]

    username = os.environ.get("GHCR_USERNAME")
    token = os.environ.get("GHCR_PULL_TOKEN")
    if not username or not token:
        click.echo(
            "No stored GHCR credential. Set GHCR_USERNAME and GHCR_PULL_TOKEN "
            "(PAT with read:packages) for the first run.",
            err=True,
        )
        sys.exit(1)
    response = http.post(
        "/containerregistryauth",
        json={"name": REGISTRY_AUTH_NAME, "username": username, "password": token},
    )
    response.raise_for_status()
    auth_id = response.json()["id"]
    click.echo(f"Registry auth created: {auth_id}")
    return auth_id


def ensure_template(http: httpx.Client, image: str, auth_id: str) -> str:
    tag = image.rsplit(":", 1)[-1]
    template_name = f"quber-parse-worker-{tag}"

    response = http.get("/templates")
    response.raise_for_status()
    for template in response.json() or []:
        if template["name"] == template_name:
            click.echo(f"Template exists: {template['id']} ({template_name})")
            return template["id"]

    response = http.post(
        "/templates",
        json={
            "name": template_name,
            "imageName": image,
            "containerRegistryAuthId": auth_id,
            "isServerless": True,
            "containerDiskInGb": 25,
            "volumeInGb": 0,
            "ports": [],
            "env": {},
        },
    )
    response.raise_for_status()
    template_id = response.json()["id"]
    click.echo(f"Template created: {template_id} ({template_name})")
    return template_id


def refresh_workers(http: httpx.Client, endpoint_id: str, workers_max: int) -> None:
    """Retire every live worker so all workers rebuild from the current template.

    Moving an endpoint to a new template does not touch its existing workers:
    FlashBoot keeps them warm on the old image and they win the race for the
    next jobs (observed live: two jobs served by a stale worker that rejected
    the new contract). Dropping workersMax to zero drains them; restoring it
    spawns fresh workers from the new template.
    """
    response = http.patch(f"/endpoints/{endpoint_id}", json={"workersMax": 0})
    response.raise_for_status()
    # A fixed pause is not enough: FlashBoot has kept workers alive through a
    # 20s drain and one of them served a job on the old image afterwards.
    # Poll until the fleet actually reports zero workers in every state.
    deadline = time.time() + WORKER_DRAIN_TIMEOUT_SECONDS
    while time.time() < deadline:
        health = httpx.get(
            f"https://api.runpod.ai/v2/{endpoint_id}/health",
            headers={"Authorization": f"Bearer {os.environ['RUNPOD_API_KEY']}"},
            timeout=30,
        ).json()
        totals = sum(health.get("workers", {}).values())
        if totals == 0:
            break
        click.echo(f"  draining... workers still present: {health['workers']}")
        time.sleep(10)
    else:
        raise click.ClickException(
            f"Workers did not drain to zero within {WORKER_DRAIN_TIMEOUT_SECONDS}s — "
            "terminate them in the console before restoring workersMax."
        )
    response = http.patch(f"/endpoints/{endpoint_id}", json={"workersMax": workers_max})
    response.raise_for_status()
    click.echo(f"Workers refreshed (drained to 0 verified, restored to {workers_max})")


def ensure_endpoint(
    http: httpx.Client, template_id: str, workers_max: int, execution_timeout_ms: int
) -> dict[str, Any]:
    response = http.get("/endpoints")
    response.raise_for_status()
    existing = next((e for e in response.json() or [] if e["name"] == ENDPOINT_NAME), None)

    converged = (
        existing is not None
        and existing.get("templateId") == template_id
        and existing.get("gpuTypeIds") == GPU_TYPE_IDS
        and existing.get("allowedCudaVersions") == ALLOWED_CUDA_VERSIONS
    )
    if existing and converged:
        click.echo(f"Endpoint exists: {existing['id']}")
        return existing
    if existing:
        # Converge template and GPU list together; either change strands
        # live workers on the old configuration, so both end in a refresh.
        response = http.patch(
            f"/endpoints/{existing['id']}",
            json={
                "templateId": template_id,
                "gpuTypeIds": GPU_TYPE_IDS,
                "allowedCudaVersions": ALLOWED_CUDA_VERSIONS,
            },
        )
        response.raise_for_status()
        moved = response.json()
        click.echo(f"Endpoint {existing['id']} moved to template {template_id}, gpus {GPU_TYPE_IDS}")
        refresh_workers(http, existing["id"], moved.get("workersMax") or workers_max)
        return moved

    response = http.post(
        "/endpoints",
        json={
            "name": ENDPOINT_NAME,
            "templateId": template_id,
            "computeType": "GPU",
            "gpuTypeIds": GPU_TYPE_IDS,
            "allowedCudaVersions": ALLOWED_CUDA_VERSIONS,
            "gpuCount": 1,
            "workersMin": 0,
            "workersMax": workers_max,
            "idleTimeout": 5,
            # FlashBoot resurrects paused workers past a drain the health
            # API reports as zero — a stale worker survived a verified-zero
            # drain and kept serving jobs on the old image. Off for
            # correctness; cold starts are acceptable at this job volume.
            "flashboot": False,
            "executionTimeoutMs": execution_timeout_ms,
        },
    )
    response.raise_for_status()
    endpoint = response.json()
    click.echo(f"Endpoint created: {endpoint['id']}")
    return endpoint


def graphql(query: str, variables: dict[str, Any] | None = None) -> dict[str, Any]:
    response = httpx.post(
        GRAPHQL_URL,
        headers={"Authorization": f"Bearer {os.environ['RUNPOD_API_KEY']}"},
        json={"query": query, "variables": variables or {}},
        timeout=60,
    )
    response.raise_for_status()
    body = response.json()
    if body.get("errors"):
        raise RuntimeError(f"GraphQL error: {body['errors']}")
    return body["data"]


def ensure_mig_exclusions(endpoint_id: str) -> None:
    """Re-assert the MIG opt-out after any endpoint create or update.

    REST writes translate gpuTypeIds into a category string and could drop
    the negative entries, so this runs last on every invocation.
    """
    data = graphql("query { myself { endpoints { id name gpuIds } } }")
    endpoint = next(e for e in data["myself"]["endpoints"] if e["id"] == endpoint_id)
    gpu_ids = endpoint["gpuIds"]
    missing = [f"-{gpu}" for gpu in MIG_EXCLUSIONS if f"-{gpu}" not in gpu_ids]
    if not missing:
        click.echo(f"MIG exclusions present: {gpu_ids}")
        return
    # The API now refuses an exclusion whose GPU type sits in a pool the
    # endpoint does not select, so each is applied on its own: the ones that
    # belong to a selected pool land, the rest are reported and skipped. An
    # exclusion for an unselected pool is moot anyway -- no worker is placed
    # there.
    for exclusion in missing:
        try:
            graphql(
                "mutation Save($input: EndpointInput!) { saveEndpoint(input: $input) { id gpuIds } }",
                {"input": {"id": endpoint_id, "name": endpoint["name"], "gpuIds": f"{gpu_ids},{exclusion}"}},
            )
        except RuntimeError as exc:
            if "does not belong to any selected GPU pool" not in str(exc):
                raise
            click.echo(f"MIG exclusion not applicable to the selected pools, skipped: {exclusion}")
            continue
        gpu_ids = f"{gpu_ids},{exclusion}"
        click.echo(f"MIG exclusion applied: {exclusion}")
    click.echo(f"gpuIds now: {gpu_ids}")


@click.command()
@click.option(
    "--image", required=True, help="Fully qualified image, e.g. ghcr.io/xmandeng/quber-runpod:v0.1.0"
)
@click.option("--workers-max", default=3, show_default=True, help="Flex worker ceiling")
@click.option(
    "--execution-timeout-ms",
    default=600_000,
    show_default=True,
    help="Per-job hard timeout. The AWS caller's deadline check keys off this value.",
)
def main(image: str, workers_max: int, execution_timeout_ms: int) -> None:
    with api_client() as http:
        auth_id = ensure_registry_auth(http)
        template_id = ensure_template(http, image, auth_id)
        endpoint = ensure_endpoint(http, template_id, workers_max, execution_timeout_ms)
    ensure_mig_exclusions(endpoint["id"])
    click.echo(f"Base URL: https://api.runpod.ai/v2/{endpoint['id']}")


if __name__ == "__main__":
    main()