#!/usr/bin/env bash
# Create (or update) the RunPod template + serverless endpoint for an endpoint
# image built by build.sh. Idempotent-ish: reuses template/endpoint by name.
#
#   ./deploy.sh qwen-edit-turbo v1
#
# Env: RUNPOD_API_KEY (required), DOCKER_HUB_USER (default plx1029),
#      GPU_IDS (comma list, default all H100 80GB variants),
#      WORKERS_MAX (default 2), IDLE_TIMEOUT (default 60s),
#      CONTAINER_DISK_GB (default 20).

set -euo pipefail
ENDPOINT="${1:?endpoint name}"; VER="${2:?version}"
USER="${DOCKER_HUB_USER:-plx1029}"
IMAGE="docker.io/${USER}/comfyui-serverless:${ENDPOINT}-${VER}"
TPL_NAME="serverless-${ENDPOINT}-${VER}"
EP_NAME="serverless-${ENDPOINT}"
GPUS="${GPU_IDS:-NVIDIA H100 80GB HBM3,NVIDIA H100 PCIe,NVIDIA H100 NVL}"
API="https://rest.runpod.io/v1"
AUTH=(-H "Authorization: Bearer $RUNPOD_API_KEY" -H "Content-Type: application/json")

gpu_json=$(python3 -c "import json,sys;print(json.dumps(sys.argv[1].split(',')))" "$GPUS")

# ---- template ----
tpl_id=$(curl -s "${AUTH[@]}" "$API/templates" | python3 -c "
import json,sys
for t in json.load(sys.stdin):
    if t.get('name')=='$TPL_NAME': print(t['id']); break")
if [ -z "$tpl_id" ]; then
    tpl_id=$(curl -s "${AUTH[@]}" -X POST "$API/templates" -d "{
        \"name\": \"$TPL_NAME\",
        \"imageName\": \"$IMAGE\",
        \"isServerless\": true,
        \"containerDiskInGb\": ${CONTAINER_DISK_GB:-20},
        \"env\": {}
    }" | python3 -c "import json,sys;d=json.load(sys.stdin);print(d.get('id') or sys.exit(json.dumps(d)))")
    echo "created template $tpl_id ($TPL_NAME -> $IMAGE)"
else
    echo "template $tpl_id ($TPL_NAME) already exists"
fi

# ---- endpoint ----
ep_id=$(curl -s "${AUTH[@]}" "$API/endpoints" | python3 -c "
import json,sys
for e in json.load(sys.stdin):
    if e.get('name')=='$EP_NAME': print(e['id']); break")
if [ -z "$ep_id" ]; then
    ep_id=$(curl -s "${AUTH[@]}" -X POST "$API/endpoints" -d "{
        \"name\": \"$EP_NAME\",
        \"templateId\": \"$tpl_id\",
        \"computeType\": \"GPU\",
        \"gpuTypeIds\": $gpu_json,
        \"gpuCount\": 1,
        \"workersMin\": 0,
        \"workersMax\": ${WORKERS_MAX:-2},
        \"idleTimeout\": ${IDLE_TIMEOUT:-60},
        \"flashboot\": true,
        \"scalerType\": \"QUEUE_DELAY\",
        \"scalerValue\": 4,
        \"executionTimeoutMs\": 600000
    }" | python3 -c "import json,sys;d=json.load(sys.stdin);print(d.get('id') or sys.exit(json.dumps(d)))")
    echo "created endpoint $ep_id ($EP_NAME)"
else
    echo "endpoint $ep_id ($EP_NAME) exists — pointing it at template $tpl_id"
    curl -s "${AUTH[@]}" -X PATCH "$API/endpoints/$ep_id" -d "{\"templateId\": \"$tpl_id\"}" >/dev/null
fi
echo
echo "endpoint id: $ep_id"
echo "  POST https://api.runpod.ai/v2/$ep_id/runsync"