#!/usr/bin/env bash
# =============================================================================
# runpod-gpu-boot.sh — arranque automático del pod GPU de video-gen
#
# Lo ejecuta el pod al iniciar (dockerStartCmd lo baja de dl.ramsesvii.com).
# Es PÚBLICO: nada de secretos acá. Lo sensible entra por variables de entorno
# del pod (GPU_AUTH_USERS = "user:bcrypt,user:bcrypt", HF_TOKEN opcional).
#
# Todo persiste en /workspace (Network Volume): venv, ComfyUI, nodos, modelos,
# binarios. El disco del contenedor se resetea en cada pod, por eso el venv
# usa --system-site-packages y reusa el torch de la imagen.
#
# Expone en :8188 un Caddy con basic_auth → ComfyUI en 127.0.0.1:8189, y
# /_gpu/phase con la fase de arranque (la lee el watchdog desde el VPS).
# Fases: booting → installing-venv → installing-comfyui → installing-deps
#        → downloading-models → starting-comfyui → ready
# =============================================================================
set -u
W=/workspace; ST=$W/status; BIN=$W/bin; COMFY=$W/ComfyUI; VENV=$W/venv
# Logs en el disco local del contenedor: un fd abierto en el volume mientras está lleno queda
# clavado en "Disk quota exceeded" aunque después haya espacio (visto el 2026-09-06 con ComfyUI).
LOGS=/var/log/videogen; mkdir -p "$LOGS"
DL="${GPU_MANIFEST_BASE:-https://dl.ramsesvii.com}"
export DEBIAN_FRONTEND=noninteractive HF_HUB_ENABLE_HF_TRANSFER=1

mkdir -p "$ST" "$BIN"
phase() { echo "$1" >"$ST/phase"; echo "[$(date -u +%FT%TZ)] phase=$1"; }
say()   { echo "[$(date -u +%FT%TZ)] $*"; }
phase booting
say "pod=${RUNPOD_POD_ID:-?} gpu=$(nvidia-smi --query-gpu=name --format=csv,noheader 2>/dev/null | head -1)"

for i in $(seq 1 30); do mountpoint -q "$W" && break; sleep 2; done
mountpoint -q "$W" || say "AVISO: /workspace no es mountpoint (¿sin Network Volume?)"

# ---- 1) Caddy: auth + fase, primero de todo para que el watchdog vea algo ----
if [ ! -x "$BIN/caddy" ]; then
  say "bajando caddy"
  curl -fsSL 'https://caddyserver.com/api/download?os=linux&arch=amd64' -o "$BIN/caddy" && chmod +x "$BIN/caddy"
fi
{
  echo "{"; echo "  admin off"; echo "}"
  echo ":8188 {"
  if [ -n "${GPU_AUTH_USERS:-}" ]; then
    echo "  basic_auth {"
    IFS=',' read -ra USERS <<<"$GPU_AUTH_USERS"
    for u in "${USERS[@]}"; do [ -n "$u" ] && echo "    ${u%%:*} ${u#*:}"; done
    echo "  }"
  fi
  echo "  handle /_gpu/* {"
  echo "    uri strip_prefix /_gpu"
  echo "    root * $ST"
  echo "    header Cache-Control no-store"
  echo "    file_server"
  echo "  }"
  echo "  handle {"
  echo "    reverse_proxy 127.0.0.1:8189 {"
  echo "      header_up -Sec-Fetch-Site"
  echo "    }"
  echo "  }"
  echo "}"
} >"$W/Caddyfile"
nohup "$BIN/caddy" run --config "$W/Caddyfile" --adapter caddyfile >"$LOGS/caddy.log" 2>&1 </dev/null &
say "caddy en :8188 (auth $( [ -n "${GPU_AUTH_USERS:-}" ] && echo on || echo OFF ))"

# ---- 2) venv persistente, atado a la versión de python+torch de la imagen ----
PYTAG="$(python3 -c 'import sys;print(f"py{sys.version_info[0]}.{sys.version_info[1]}")')-${PYTORCH_VERSION:-torch}"
if [ ! -x "$VENV/bin/python" ] || [ "$(cat "$ST/venv-tag" 2>/dev/null)" != "$PYTAG" ]; then
  phase installing-venv
  rm -rf "$VENV"
  python3 -m venv --system-site-packages "$VENV" 2>/dev/null \
    || { say "venv no disponible, uso virtualenv"; python3 -m pip install -q virtualenv && python3 -m virtualenv -q --system-site-packages "$VENV"; }
  [ -x "$VENV/bin/python" ] || { say "ERROR: no pude crear el venv"; phase failed-venv; exit 1; }
  echo "$PYTAG" >"$ST/venv-tag"; rm -f "$ST/deps-ok"
fi
PY="$VENV/bin/python"; PIP="$PY -m pip"
$PIP install -q --upgrade pip >/dev/null 2>&1

# ---- 3) ComfyUI + custom nodes (manifiesto: runpod-gpu-nodes.txt) ----
if [ ! -d "$COMFY/.git" ]; then
  phase installing-comfyui
  git clone --depth 1 https://github.com/comfyanonymous/ComfyUI "$COMFY"; rm -f "$ST/deps-ok"
fi
NODES=$(curl -fsSL "$DL/runpod-gpu-nodes.txt" 2>/dev/null || true)
if [ -n "$NODES" ]; then
  while IFS= read -r url; do
    url="${url%%#*}"; url="$(echo "$url" | xargs)"; [ -z "$url" ] && continue
    name="$(basename "$url" .git)"; d="$COMFY/custom_nodes/$name"
    if [ ! -d "$d/.git" ]; then phase installing-comfyui; say "nodo $name"; git clone --depth 1 "$url" "$d" && rm -f "$ST/deps-ok"; fi
  done <<<"$NODES"
fi

# ---- 4) deps de python (una vez por venv / por nodo nuevo) ----
if [ ! -f "$ST/deps-ok" ]; then
  phase installing-deps
  $PIP install -q -r "$COMFY/requirements.txt"
  for r in "$COMFY"/custom_nodes/*/requirements.txt; do [ -f "$r" ] && $PIP install -q -r "$r"; done
  $PIP install -q "huggingface_hub[cli,hf_transfer]"
  touch "$ST/deps-ok"
fi

# ---- 5) modelos: delegado a models-sync.sh (mismo manifiesto; loguea fallos en models-sync.log) ----
curl -fsSL "$DL/runpod-gpu-models-sync.sh" | HF_HOME="$W/hf-cache" bash || say "models-sync terminó con fallos (ver /workspace/models-sync.log)"

# ---- 6) workflows (manifiesto: runpod-gpu-workflows.txt  url[|nombre]) ----
WFS=$(curl -fsSL "$DL/runpod-gpu-workflows.txt" 2>/dev/null || true)
if [ -n "$WFS" ]; then
  mkdir -p "$COMFY/user/default/workflows"
  while IFS='|' read -r url name; do
    url="$(echo "${url%%#*}" | xargs)"; [ -z "$url" ] && continue
    name="$(echo "${name:-}" | xargs)"; [ -n "$name" ] || name="$(basename "$url")"
    curl -fsSL "$url" -o "$COMFY/user/default/workflows/$name" 2>/dev/null || say "FALLÓ workflow $name"
  done <<<"$WFS"
fi

# ---- 7) ComfyUI ----
phase starting-comfyui
cd "$COMFY" || exit 1
nohup "$PY" main.py --listen 127.0.0.1 --port 8189 >"$LOGS/comfyui.log" 2>&1 </dev/null &
for i in $(seq 1 120); do
  curl -s -o /dev/null http://127.0.0.1:8189/system_stats && { phase ready; say "ComfyUI listo"; exit 0; }
  sleep 5
done
phase failed-comfyui
say "ComfyUI no respondió en 10 min; ver $LOGS/comfyui.log"