#!/usr/bin/env bash
#
# /post_start.sh — comfyui-mcp boot hook (FAST-RESTART design).
# =============================================================================
# The base image's CMD is /start.sh (from runpod/containers). It:
#   1. service nginx start          (reads our /etc/nginx/nginx.conf: :3000->:3001)
#   2. runs /pre_start.sh (if any)
#   3. setup_ssh (RunPod $PUBLIC_KEY injection) + start_jupyter ($JUPYTER_PASSWORD)
#   4. runs THIS /post_start.sh
#   5. sleep infinity                (keeps the pod alive)
#
# So by the time we run, nginx + sshd + JupyterLab are already up.
#
# FAST RESTART: the ComfyUI install + venv + custom_nodes are BAKED IN THE IMAGE
# at ${COMFY_HOME} (default /opt/ComfyUI) and run DIRECTLY from there. We do NOT
# seed/sync/rsync anything onto /workspace. The ONLY volume prep is a fast,
# idempotent `mkdir -p` of the user-data dirs. ComfyUI is then pointed at those
# dirs via per-directory launch flags + extra_model_paths.yaml. A warm restart
# therefore does NO install/sync/seed — just mkdir + launch (~30-60s init).
#
# WHAT PERSISTS / WHAT DOESN'T:
#   * PERSIST (on /workspace): user/ (workflows+settings+Manager config),
#     models/ (incl. Manager downloads), input/, output/, and — since §4.5 below —
#     custom_nodes/ too (symlinked onto the volume; nodes the agent/Manager
#     install at runtime NOW SURVIVE a restart). The image's baked nodes (incl.
#     the Agent Panel) are seeded/refreshed onto the volume every boot without
#     clobbering nodes the user installed themselves.
#   * EPHEMERAL (in the container): the ComfyUI install + venv + all caches —
#     still baked/fast per the fast-restart design above.
#
# IMPORTANT: this script must end alive-and-non-fatal. The base runs it with
# `set -e`, so if we returned non-zero the pod would die. We launch ComfyUI in
# the background and `exec tail -F` its log: that streams ComfyUI's output to the
# RunPod console AND holds the process open (== keeps the pod up) regardless of
# whether ComfyUI later crashes.
# =============================================================================
set -uo pipefail   # NOT -e: services are best-effort; we must stay alive.

log() { echo "[comfyui-mcp/post_start] $*"; }

# ---- Config (override via pod Environment) ----------------------------------
COMFY_HOME="${COMFY_HOME:-/opt/ComfyUI}"               # BAKED ComfyUI (image)
SEED_MODELS="${SEED_MODELS:-/opt/ComfyUI-seed-models}" # baked spotcheck model(s)
WORKSPACE="${WORKSPACE:-/workspace}"                   # network volume (USER DATA)
COMFY_PORT="${COMFY_PORT:-3001}"                        # nginx :3000 -> here
COMFY_NETWORK_MODE="${COMFY_NETWORK_MODE:-personal_cloud}"
# "weak" so the Agent Panel can install custom nodes from ARBITRARY git URLs
# (Manager classes those high-risk and silently skips them at "normal-", while
# registry-id installs pass — a confusing half-working state). This is a
# single-user pod whose entire premise is agent-driven installs; set
# COMFY_SECURITY_LEVEL=normal- to restore the guardrails.
COMFY_SECURITY_LEVEL="${COMFY_SECURITY_LEVEL:-weak}"
COMFY_EXTRA_ARGS="${COMFY_EXTRA_ARGS:-}"              # extra ComfyUI flags
EXTRA_MODEL_PATHS="${EXTRA_MODEL_PATHS:-${COMFY_HOME}/extra_model_paths.yaml}"
# Pull the latest Agent Panel release on every boot (git fetch + reset --hard,
# BEFORE ComfyUI launches — see §4.5b) instead of waiting for a new pod image.
# Automatically a no-op for a PANEL_REF-pinned build (detached HEAD — pinning
# means the user wants reproducibility, not drift). Set to 0 to disable outright.
PANEL_AUTO_UPDATE="${PANEL_AUTO_UPDATE:-1}"
# Where the §4.5(b.2) self-heal re-clones the Agent Panel from if the copy on the
# volume is broken (0-byte/missing files). Baked as ENV by the Dockerfile;
# PANEL_REF (optional tag/branch) keeps a ref-pinned build pinned through a heal.
PANEL_REPO="${PANEL_REPO:-https://github.com/artokun/comfyui-mcp-panel.git}"
PANEL_REF="${PANEL_REF:-}"

# Volume user-data dirs (the ONLY things on /workspace).
USER_DIR="${WORKSPACE}/user"
MODELS_DIR="${WORKSPACE}/models"
INPUT_DIR="${WORKSPACE}/input"
OUTPUT_DIR="${WORKSPACE}/output"

# Logs go to the EPHEMERAL container fs, NOT the volume. The RunPod console
# streams them live (we `exec tail -F` below), and the fast-restart contract
# keeps /workspace EXACTLY user/models/input/output — logs are runtime cruft.
LOG_DIR="${COMFY_LOG_DIR:-/var/log/comfyui-mcp}"
mkdir -p "${LOG_DIR}"

# Minimum host NVIDIA driver. Baked per image variant by the Dockerfile
# (MIN_DRIVER_DEFAULT env: 570 for the default cu128 build — the same CUDA 12.8
# bar the runpod/pytorch base's container-start gate already enforces — 580 for
# the cu130 perf variant, which needs CUDA 13). Override at runtime with
# MIN_DRIVER, or set MIN_DRIVER=0 to disable the check.
MIN_DRIVER="${MIN_DRIVER:-${MIN_DRIVER_DEFAULT:-570}}"

# -----------------------------------------------------------------------------
# 0. GPU DRIVER PREFLIGHT. The host NVIDIA driver is NOT upgradable from inside the
#    container, so a too-old driver (common when the scheduler drops a new GPU on a
#    stale host) makes torch fail to init CUDA and ComfyUI crash-loops with a cryptic
#    error. Detect it UP FRONT (before launching ComfyUI) and hold the pod open with a
#    clear "redeploy on a newer host" message instead of looping.
# -----------------------------------------------------------------------------
if [ "${MIN_DRIVER}" != "0" ] && command -v nvidia-smi >/dev/null 2>&1; then
  GPU_NAME="$(nvidia-smi --query-gpu=name --format=csv,noheader 2>/dev/null | head -1)"
  DRV="$(nvidia-smi --query-gpu=driver_version --format=csv,noheader 2>/dev/null | head -1)"
  DRV_MAJOR="${DRV%%.*}"
  log "GPU: ${GPU_NAME:-unknown} | driver: ${DRV:-unknown} (this image needs driver >= ${MIN_DRIVER})"
  if [ -n "${DRV_MAJOR}" ] && [ "${DRV_MAJOR}" -lt "${MIN_DRIVER}" ] 2>/dev/null; then
    log "============================================================================"
    log "FATAL: host NVIDIA driver ${DRV} is TOO OLD for this image (needs >= ${MIN_DRIVER})."
    log "  Your ${GPU_NAME:-GPU} (esp. RTX 50xx / Blackwell) needs a newer driver, and the"
    log "  HOST driver CANNOT be upgraded from inside a pod — torch would fail to init CUDA."
    log "  FIX: TERMINATE this pod and REDEPLOY on a host with driver >= ${MIN_DRIVER}"
    log "       (filter by 'CUDA Version' on the RunPod deploy screen). Verify with: nvidia-smi"
    log "  Holding the pod open (no crash-loop) so you can inspect it. Set MIN_DRIVER=0 to bypass."
    log "============================================================================"
    exec sleep infinity
  fi
else
  [ "${MIN_DRIVER}" = "0" ] && log "driver preflight disabled (MIN_DRIVER=0)." \
    || log "WARNING: nvidia-smi not found — skipping GPU driver preflight."
fi

# -----------------------------------------------------------------------------
# 1. Volume prep — the ONLY boot-time volume work. Fast + idempotent.
#    Create user/models/input/output + the model category subfolders that
#    extra_model_paths.yaml maps, so the dirs exist on a cold volume (ComfyUI /
#    Manager would create them lazily, but pre-creating keeps the UI tidy and
#    guarantees Manager has a place to download into).
#    NO caches, NO venv, NO custom_nodes are placed on the volume.
# -----------------------------------------------------------------------------
# Free-space preflight. A FULL volume is the nastiest failure mode on this pod:
# every `cp`/`git checkout` still CREATES its files but writes 0 bytes into them
# (ENOSPC), so the panel + custom nodes turn into a correct-looking tree of empty
# husks that PERSISTS across redeploys. Detect it up front and say so in plain
# words instead of letting the boot "succeed".
free_mb() { df -Pm "$1" 2>/dev/null | awk 'NR==2 {print $4}'; }
WS_FREE_MB="$(free_mb "${WORKSPACE}")"
log "volume free space: ${WS_FREE_MB:-unknown} MB on ${WORKSPACE}"
if [ -n "${WS_FREE_MB}" ] && [ "${WS_FREE_MB}" -lt 500 ] 2>/dev/null; then
  log "============================================================================"
  log "WARNING: ${WORKSPACE} has only ${WS_FREE_MB} MB free. Writes will fail with"
  log "  ENOSPC and leave 0-BYTE files (empty panel, broken custom nodes, failed"
  log "  model downloads). Free up space or grow the network volume, then restart."
  log "============================================================================"
fi

log "preparing /workspace user-data dirs (mkdir -p; no sync)…"
mkdir -p "${USER_DIR}" "${INPUT_DIR}" "${OUTPUT_DIR}"
for sub in checkpoints configs loras vae text_encoders clip diffusion_models \
           unet clip_vision style_models embeddings diffusers vae_approx \
           controlnet t2i_adapter gligen upscale_models latent_upscale_models \
           hypernetworks photomaker classifiers model_patches audio_encoders \
           background_removal frame_interpolation geometry_estimation \
           optical_flow detection \
           sams ultralytics ultralytics/bbox ultralytics/segm; do
  mkdir -p "${MODELS_DIR}/${sub}"
done

# -----------------------------------------------------------------------------
# 2. Spotcheck model — first-boot only. If the SDXL checkpoint isn't on the
#    volume yet and a baked copy exists, copy it so it's visible AND persists.
#    (BAKE_SPOTCHECK_MODEL=0 builds omit the baked copy → this is a no-op.)
# -----------------------------------------------------------------------------
if [ ! -f "${MODELS_DIR}/checkpoints/sd_xl_base_1.0.safetensors" ] \
   && ls "${SEED_MODELS}"/*.safetensors >/dev/null 2>&1; then
  # Guard: the copy is ~7 GB. On a small/nearly-full volume it fills the disk and
  # everything AFTER it (panel seed, node installs) degrades into 0-byte files.
  # The spotcheck model is a convenience — skip it rather than poison the volume.
  SEED_MODELS_MB="$(du -sm "${SEED_MODELS}" 2>/dev/null | awk '{print $1}')"
  WS_FREE_MB="$(free_mb "${WORKSPACE}")"
  if [ -n "${WS_FREE_MB}" ] && [ -n "${SEED_MODELS_MB}" ] \
     && [ "${WS_FREE_MB}" -lt "$((SEED_MODELS_MB + 2048))" ] 2>/dev/null; then
    log "SKIP spotcheck model: needs ~${SEED_MODELS_MB} MB (+2 GB headroom) but only ${WS_FREE_MB} MB free on ${WORKSPACE}."
  else
    log "first boot: copying baked spotcheck model(s) into models/checkpoints…"
    cp -n "${SEED_MODELS}"/*.safetensors "${MODELS_DIR}/checkpoints/" \
      && log "spotcheck model in place." \
      || log "WARN: spotcheck copy failed (continuing)."
  fi
fi

# -----------------------------------------------------------------------------
# 3. Manager remote-install gate. The RUNNING ComfyUI reads its Manager config
#    UNDER the --user-directory (= ${USER_DIR}). Depending on the ComfyUI build,
#    comfyui_manager looks at either:
#        ${USER_DIR}/__manager/config.ini              (new: System User API)
#        ${USER_DIR}/default/ComfyUI-Manager/config.ini (older)
#    We write/re-assert BOTH on every boot so the gate is correct regardless.
#    network_mode=personal_cloud + a permissive security_level are REQUIRED for
#    the /v2 install-model gate the Agent Panel uses.
# -----------------------------------------------------------------------------
write_manager_config() {  # $1 = config.ini absolute path
  local cfg="$1" dir
  dir="$(dirname "${cfg}")"
  mkdir -p "${dir}"
  if [ ! -f "${cfg}" ]; then
    # Prefer the baked template if present, else synthesize.
    if [ -f "${COMFY_HOME}/config.ini.seed" ]; then
      cp "${COMFY_HOME}/config.ini.seed" "${cfg}"
    else
      printf '[default]\nnetwork_mode = %s\nsecurity_level = %s\n' \
        "${COMFY_NETWORK_MODE}" "${COMFY_SECURITY_LEVEL}" > "${cfg}"
    fi
  fi
  # Re-assert the two keys from env (idempotent; survives template drift).
  if grep -q '^network_mode' "${cfg}"; then
    sed -i "s/^network_mode.*/network_mode = ${COMFY_NETWORK_MODE}/" "${cfg}"
  else
    printf '\nnetwork_mode = %s\n' "${COMFY_NETWORK_MODE}" >> "${cfg}"
  fi
  if grep -q '^security_level' "${cfg}"; then
    sed -i "s/^security_level.*/security_level = ${COMFY_SECURITY_LEVEL}/" "${cfg}"
  else
    printf 'security_level = %s\n' "${COMFY_SECURITY_LEVEL}" >> "${cfg}"
  fi
}
write_manager_config "${USER_DIR}/__manager/config.ini"
write_manager_config "${USER_DIR}/default/ComfyUI-Manager/config.ini"
log "Manager config asserted: network_mode=${COMFY_NETWORK_MODE} security_level=${COMFY_SECURITY_LEVEL}"

# -----------------------------------------------------------------------------
# 3.5 Manager AUTO-UPDATE (honors the template's COMFY_AUTOUPDATE_MANAGER env,
#     which existed on the template for ages but was silently ignored). The
#     venv is EPHEMERAL by design, so any Manager fix applied inside a running
#     pod (e.g. "update to nightly" to cure a broken install path) EVAPORATES
#     on the next stop/start — this is the only runtime-persistent channel for
#     Manager fixes short of shipping a new image.
#       COMFY_AUTOUPDATE_MANAGER=1        (default) latest STABLE pip release
#       COMFY_AUTOUPDATE_MANAGER=nightly  ComfyUI-Manager git main (bleeding edge)
#       COMFY_AUTOUPDATE_MANAGER=0        keep the version baked in the image
#     Best-effort with a hard timeout — an offline pod or a PyPI hiccup must
#     never block boot. The baked manager_core shim is a SEPARATE site-packages
#     module, so upgrades leave it in place (it goes inert if a future Manager
#     drops the legacy import).
# -----------------------------------------------------------------------------
COMFY_AUTOUPDATE_MANAGER="${COMFY_AUTOUPDATE_MANAGER:-1}"
MGR_PIP="${COMFY_HOME}/venv/bin/pip"
MGR_PY="${COMFY_HOME}/venv/bin/python"
mgr_ver() { "${MGR_PY}" -c "import importlib.metadata as m; print(m.version('comfyui-manager'))" 2>/dev/null || echo unknown; }
if [ -x "${MGR_PIP}" ] && [ "${COMFY_AUTOUPDATE_MANAGER}" != "0" ]; then
  MGR_BEFORE="$(mgr_ver)"
  case "${COMFY_AUTOUPDATE_MANAGER}" in
    nightly) MGR_SPEC="git+https://github.com/Comfy-Org/ComfyUI-Manager.git@main" ;;
    *)       MGR_SPEC="comfyui_manager" ;;
  esac
  if timeout 180 "${MGR_PIP}" install -q -U --retries 2 "${MGR_SPEC}" \
       >>"${LOG_DIR}/manager-update.log" 2>&1; then
    MGR_AFTER="$(mgr_ver)"
    if [ "${MGR_BEFORE}" = "${MGR_AFTER}" ]; then
      log "Manager up to date: ${MGR_AFTER} (COMFY_AUTOUPDATE_MANAGER=${COMFY_AUTOUPDATE_MANAGER})"
    else
      log "Manager auto-updated: ${MGR_BEFORE} -> ${MGR_AFTER} (COMFY_AUTOUPDATE_MANAGER=${COMFY_AUTOUPDATE_MANAGER})"
    fi
  else
    log "WARN: Manager auto-update failed/timed out — keeping baked $(mgr_ver) (see manager-update.log)"
  fi
else
  [ "${COMFY_AUTOUPDATE_MANAGER}" = "0" ] \
    && log "Manager auto-update disabled (COMFY_AUTOUPDATE_MANAGER=0) — baked $(mgr_ver)"
fi

# -----------------------------------------------------------------------------
# 4. Ancillary services (best-effort — skipped if the binary is absent).
# -----------------------------------------------------------------------------
service cron start >/dev/null 2>&1 && log "cron started" || log "cron not available (skip)"

if command -v code-server >/dev/null 2>&1; then
  log "starting code-server on :8080 (nginx front :8081)…"
  nohup code-server --bind-addr 0.0.0.0:8080 --auth none \
    >"${LOG_DIR}/code-server.log" 2>&1 &
else
  log "code-server not installed (skip)"
fi

if command -v runpod-uploader >/dev/null 2>&1; then
  log "starting runpod-uploader…"
  nohup runpod-uploader >"${LOG_DIR}/runpod-uploader.log" 2>&1 &
else
  log "runpod-uploader not present (skip)"
fi

if [ -f /app-manager/app.js ] && command -v node >/dev/null 2>&1; then
  log "starting app-manager on :8000 (nginx front :8001)…"
  ( cd /app-manager && nohup node app.js >"${LOG_DIR}/app-manager.log" 2>&1 & )
else
  log "app-manager not present or node missing (skip)"
fi

command -v croc >/dev/null 2>&1 && log "croc available (on-demand P2P transfer)"

# Dead-man switch (#269): pods created via the runpod tool's create action carry a heartbeat
# server (:8189, token-gated) + a watchdog loop that STOPS THE POD if
# comfyui-mcp's heartbeats stop (orchestrator crash/offline = the in-process
# idle auto-stop died with it). The stop is authorized by the POD-SCOPED
# RUNPOD_API_KEY RunPod auto-injects (never the owner's account key). Arms
# only when the pod env carries DEADMAN_TOKEN — console-deployed pods stay inert.
if [ "${DEADMAN_DISABLE:-0}" = "1" ]; then
  log "dead-man switch disabled (DEADMAN_DISABLE=1)"
elif [ -n "${RUNPOD_API_KEY:-}" ] && [ -n "${RUNPOD_POD_ID:-}" ] && [ -n "${DEADMAN_TOKEN:-}" ] \
     && [ -f /opt/comfyui-mcp-deadman/deadman_server.py ] && [ -x /opt/comfyui-mcp-deadman/deadman_watch.sh ]; then
  log "starting dead-man heartbeat server (:8189) + watchdog (boot grace ${DEADMAN_BOOT_GRACE_S:-2700}s, beat grace ${DEADMAN_BEAT_GRACE_S:-1200}s)…"
  nohup "${COMFY_HOME}/venv/bin/python" /opt/comfyui-mcp-deadman/deadman_server.py >>"${LOG_DIR}/deadman.log" 2>&1 &
  LOG_DIR="${LOG_DIR}" nohup /opt/comfyui-mcp-deadman/deadman_watch.sh >>"${LOG_DIR}/deadman.log" 2>&1 &
else
  log "dead-man switch inert (no RUNPOD_API_KEY/DEADMAN_TOKEN on the pod — only pods deployed by the runpod tool's create action carry it)"
fi

# File Browser — web file manager for /workspace (browse/upload/download/delete),
# fronted by nginx on :8083. noauth (parity with code-server --auth none — the
# RunPod proxy URL is the boundary; set FILEBROWSER_PASSWORD for a login). The DB
# is EPHEMERAL (re-init each boot) so it never clutters /workspace.
if command -v filebrowser >/dev/null 2>&1; then
  FB_DB=/var/lib/comfyui-mcp/filebrowser.db
  mkdir -p /var/lib/comfyui-mcp
  rm -f "${FB_DB}"                       # fresh DB each boot (ephemeral, idempotent)
  filebrowser -d "${FB_DB}" config init >>"${LOG_DIR}/filebrowser.log" 2>&1
  filebrowser -d "${FB_DB}" config set --root /workspace >>"${LOG_DIR}/filebrowser.log" 2>&1
  if [ -n "${FILEBROWSER_PASSWORD:-}" ]; then
    filebrowser -d "${FB_DB}" config set --auth.method=json >>"${LOG_DIR}/filebrowser.log" 2>&1
    filebrowser -d "${FB_DB}" users add admin "${FILEBROWSER_PASSWORD}" --perm.admin \
      >>"${LOG_DIR}/filebrowser.log" 2>&1 || true
    log "starting filebrowser on :8082 (nginx front :8083; login admin/\$FILEBROWSER_PASSWORD)…"
  else
    filebrowser -d "${FB_DB}" config set --auth.method=noauth >>"${LOG_DIR}/filebrowser.log" 2>&1
    log "starting filebrowser on :8082 (nginx front :8083; NO auth — set FILEBROWSER_PASSWORD to lock)…"
  fi
  nohup filebrowser -d "${FB_DB}" -r /workspace -a 0.0.0.0 -p 8082 \
    >>"${LOG_DIR}/filebrowser.log" 2>&1 &
else
  log "filebrowser not installed (skip)"
fi

# HuggingFace auth for gated downloads: huggingface_hub (used by ComfyUI + many
# custom nodes) reads HF_TOKEN. Accept our MCP's HUGGINGFACE_TOKEN name as an
# alias so setting EITHER on the pod works. Exported here so the ComfyUI launch
# below (and Manager) inherit it.
export HF_TOKEN="${HF_TOKEN:-${HUGGINGFACE_TOKEN:-}}"
[ -n "${HF_TOKEN}" ] && log "HF_TOKEN present — gated HuggingFace downloads authenticated."

# -----------------------------------------------------------------------------
# 4.4 FAST DOWNLOADS — two independent paths:
#     (a) huggingface_hub fetches (custom nodes, ComfyUI internals): hf_transfer
#         (Rust, parallel) + xet high-performance mode, DEFAULT ON for pod use
#         (datacenter pipe, no proxies) — set either var to 0 on the pod to opt
#         out (that is also the first troubleshooting step for a failing HF
#         download: hf_transfer trades resume robustness for speed).
#     (b) Manager install-model: an aria2 RPC sidecar. Manager's built-in
#         downloader (single stream, tiny chunks) collapses to <1-4 MB/s
#         against the MooseFS network volume, wedging the serial install queue
#         for hours per model; with COMFYUI_MANAGER_ARIA2_SERVER set, Manager
#         hands downloads to local aria2 instead (multi-connection → full pipe
#         speed; the same pod measured 792 kB/s built-in vs 15-80 MB/s for
#         aria2-class fetches). Set ARIA2_DISABLE=1 to opt out.
#     GUARDS (review findings, 2026-07-08):
#       - the flag/env is only set when the matching package imports — both the
#         hub (HF_HUB_ENABLE_HF_TRANSFER) and Manager (import aria2p at startup
#         when COMFYUI_MANAGER_ARIA2_SERVER is set) RAISE/CRASH otherwise;
#       - aria2 runs under a respawn loop (a --daemon'd process that crashes
#         mid-session would leave Manager hard-failing with NO fallback);
#       - the Manager env is exported ONLY after a real RPC answer (daemonizing
#         successfully says nothing about the listener being ready);
#       - /models → volume symlink: Manager's aria2 path joins RELATIVE model
#         dirs onto '/models' (container-ephemeral!) — the symlink turns that
#         silent-loss edge case into a persistent write;
#       - the RPC secret lives in a 600 conf file (not the process cmdline) and
#         is length-checked (an empty --rpc-secret must never ship).
# -----------------------------------------------------------------------------
if "${COMFY_HOME}/venv/bin/python" -c "import hf_transfer" >/dev/null 2>&1; then
  export HF_HUB_ENABLE_HF_TRANSFER="${HF_HUB_ENABLE_HF_TRANSFER:-1}"
  export HF_XET_HIGH_PERFORMANCE="${HF_XET_HIGH_PERFORMANCE:-1}"
  log "HF fast downloads: HF_HUB_ENABLE_HF_TRANSFER=${HF_HUB_ENABLE_HF_TRANSFER} HF_XET_HIGH_PERFORMANCE=${HF_XET_HIGH_PERFORMANCE} (set 0 on the pod to opt out)"
else
  log "hf_transfer not importable — HF hub downloads use the default backend"
fi

ARIA2_RPC_PORT="${ARIA2_RPC_PORT:-6800}"
RUNSTATE_DIR=/var/lib/comfyui-mcp
if [ "${ARIA2_DISABLE:-0}" = "1" ]; then
  log "aria2 sidecar disabled (ARIA2_DISABLE=1) — Manager uses its built-in downloader"
elif command -v aria2c >/dev/null 2>&1 \
     && "${COMFY_HOME}/venv/bin/python" -c "import aria2p" >/dev/null 2>&1; then
  ARIA2_RPC_SECRET="${ARIA2_RPC_SECRET:-$(od -An -N16 -tx1 /dev/urandom 2>/dev/null | tr -d ' \n')}"
  if [ "${#ARIA2_RPC_SECRET}" -lt 8 ]; then
    log "WARN: could not generate an aria2 RPC secret — sidecar skipped (built-in downloader)"
  else
    mkdir -p "${RUNSTATE_DIR}"
    ARIA2_CONF="${RUNSTATE_DIR}/aria2.conf"
    ( umask 077; printf 'rpc-secret=%s\n' "${ARIA2_RPC_SECRET}" > "${ARIA2_CONF}" )
    # /models guard — see GUARDS above. ln, not mkdir: an existing real /models
    # (never shipped by this image) is left alone.
    [ -e /models ] || ln -s "${MODELS_DIR}" /models 2>/dev/null \
      || log "WARN: could not link /models -> ${MODELS_DIR}"
    (
      while :; do
        aria2c --conf-path="${ARIA2_CONF}" --enable-rpc --rpc-listen-all=false \
          --rpc-listen-port="${ARIA2_RPC_PORT}" \
          --max-connection-per-server=16 --split=16 --min-split-size=8M \
          --file-allocation=none --continue=true \
          --allow-overwrite=true --auto-file-renaming=false \
          --console-log-level=warn >>"${LOG_DIR}/aria2.log" 2>&1
        echo "[comfyui-mcp/post_start] aria2c exited (rc=$?) — restarting in 3s" >>"${LOG_DIR}/aria2.log"
        sleep 3
      done
    ) &
    ARIA2_SUPERVISOR_PID=$!
    ARIA2_READY=0
    for _ in $(seq 1 40); do
      if curl -s -m 2 "http://127.0.0.1:${ARIA2_RPC_PORT}/jsonrpc" \
           -H 'Content-Type: application/json' \
           -d "{\"jsonrpc\":\"2.0\",\"id\":\"boot\",\"method\":\"aria2.getVersion\",\"params\":[\"token:${ARIA2_RPC_SECRET}\"]}" \
           2>/dev/null | grep -q '"result"'; then
        ARIA2_READY=1
        break
      fi
      sleep 0.5
    done
    if [ "${ARIA2_READY}" = "1" ]; then
      export COMFYUI_MANAGER_ARIA2_SERVER="http://127.0.0.1:${ARIA2_RPC_PORT}"
      export COMFYUI_MANAGER_ARIA2_SECRET="${ARIA2_RPC_SECRET}"
      # Root-only state file so the boot test (and a debugging human) can make a
      # REAL RPC call with the live secret instead of trusting a log line.
      ( umask 077; printf 'COMFYUI_MANAGER_ARIA2_SERVER=%s\nCOMFYUI_MANAGER_ARIA2_SECRET=%s\n' \
          "${COMFYUI_MANAGER_ARIA2_SERVER}" "${ARIA2_RPC_SECRET}" > "${RUNSTATE_DIR}/aria2-rpc.env" )
      log "aria2 RPC sidecar up (127.0.0.1:${ARIA2_RPC_PORT}, RPC-verified, supervised) — Manager model downloads now multi-connection"
    else
      kill "${ARIA2_SUPERVISOR_PID}" 2>/dev/null || true
      pkill -x aria2c 2>/dev/null || true
      log "WARN: aria2 RPC did not answer within 20s — sidecar disabled, Manager falls back to its built-in downloader"
    fi
  fi
else
  log "aria2c/aria2p not baked in this image — Manager uses its built-in downloader"
fi

# -----------------------------------------------------------------------------
# 4.5 CUSTOM NODES → the VOLUME (so runtime-installed nodes SURVIVE a restart).
#     The base ComfyUI + venv stay in the image (fast boot), but custom_nodes live
#     on /workspace — the #1 user complaint with a pure image-baked custom_nodes
#     was losing them on stop/start. Mechanism:
#       (a) symlink ${COMFY_HOME}/custom_nodes -> /workspace/custom_nodes, so
#           ComfyUI AND Manager read/write the volume with zero path changes;
#       (b) seed/refresh the image's baked nodes (panel + ComfyUI builtins) into it
#           every boot — image upgrades push a fresh panel while USER nodes persist;
#       (b.1) fast-forward the Agent Panel specifically to its latest release via
#           git, independent of the baked image (see below) — so a panel release
#           reaches pods without waiting for a new image build;
#       (c) reinstall each node's Python deps into the (ephemeral, image) venv from
#           a PERSISTENT pip cache on the volume — required because the venv is in
#           the image: node CODE persists on the volume, its DEPS must be
#           re-materialized into the venv each boot (fast after the first time).
# -----------------------------------------------------------------------------
CN_VOL="${WORKSPACE}/custom_nodes"
CN_LINK="${COMFY_HOME}/custom_nodes"
CN_SEED="${COMFY_HOME}/custom_nodes_seed"
export PIP_CACHE_DIR="${WORKSPACE}/.cache/pip"
# Pack install scripts resolve their OWN download dir from $COMFYUI_MODEL_PATH and
# fall back to the IMAGE's models dir when it is unset. extra_model_paths.yaml cannot
# reach them: it adds SEARCH paths, it does not move folder_paths.models_dir. So
# Impact Subpack's installer drops face_yolov8m.pt into
# /opt/ComfyUI/models/ultralytics/bbox (it has no folder_paths fallback at all) and
# Impact Pack's drops sam_vit_b_01ec64.pth into /opt/ComfyUI/models/sams - both on the
# ephemeral layer, both gone on the next rebuild (#2302). ComfyUI core never reads this
# variable (models_dir comes from --models-directory or base_path), so setting it
# redirects ONLY the node packs' own installers, and onto the same dirs
# extra_model_paths.yaml already maps.
export COMFYUI_MODEL_PATH="${MODELS_DIR}"
mkdir -p "${CN_VOL}" "${PIP_CACHE_DIR}"

# (a) Point the image's custom_nodes at the volume. Replace whatever is there — a
#     real dir on a fresh container, or a stale symlink — with a link to the vol.
if [ ! -L "${CN_LINK}" ] || [ "$(readlink -f "${CN_LINK}")" != "$(readlink -f "${CN_VOL}")" ]; then
  rm -rf "${CN_LINK}"
  ln -s "${CN_VOL}" "${CN_LINK}"
  log "custom_nodes -> ${CN_VOL} (symlinked; installs now persist on the volume)"
fi

# (b) Seed/refresh the baked nodes (panel + builtins) onto the volume. cp -rf
#     overwrites the image-owned copies (keeps them current on an image upgrade)
#     but never deletes the user's OWN nodes already on the volume. Errors go to
#     a LOG, not /dev/null — a swallowed ENOSPC here once shipped a panel whose
#     every file existed with 0 bytes, and nothing in the console said why.
if [ -d "${CN_SEED}" ]; then
  if cp -rf "${CN_SEED}/." "${CN_VOL}/" 2>>"${LOG_DIR}/custom-nodes-seed.log"; then
    log "seeded/refreshed baked custom_nodes (panel + builtins) onto the volume"
  else
    log "WARN: custom_nodes seed refresh had errors — first lines:"
    head -3 "${LOG_DIR}/custom-nodes-seed.log" 2>/dev/null | while IFS= read -r l; do log "  cp: ${l}"; done
    log "  (full log: ${LOG_DIR}/custom-nodes-seed.log; free space: $(free_mb "${WORKSPACE}") MB)"
  fi
fi

# (b.1) BAKED-NODE AUTO-UPDATE — decouples "get the latest Agent Panel /
#     Crystools" from "wait for a new pod image". (a)+(b) above only refresh the
#     volume from what THIS IMAGE baked at build time — which also means (b)
#     REVERTS any newer version a user pulled via Manager back to the baked one
#     on every boot. Every baked node is a plain git clone (see the Dockerfile),
#     so fast-forward each in place before ComfyUI launches (no extra restart
#     needed, unlike Manager's update API which applies on the NEXT launch).
#
#     Per-node skip rules:
#       * detached HEAD (a build pinned PANEL_REF/CRYSTOOLS_REF for
#         reproducibility) — never silently overridden;
#       * PANEL_AUTO_UPDATE!=1 disables the lot (same knob as before).
#     Best-effort: any failure (offline pod, rate limit, force-pushed history)
#     logs a warning and keeps the existing checkout.
ff_baked_node() {  # $1 = dir under custom_nodes, $2 = display name
  local dir="${CN_VOL}/$1" label="$2" branch
  [ -d "${dir}/.git" ] || return 0
  branch="$(git -C "${dir}" symbolic-ref -q --short HEAD 2>/dev/null || true)"
  if [ -z "${branch}" ]; then
    log "${label} checkout is pinned (detached HEAD) — skipping auto-update, as intended"
    return 0
  fi
  log "checking for a newer ${label} (branch: ${branch})…"
  if git -C "${dir}" fetch --depth 1 origin "${branch}" >>"${LOG_DIR}/panel-update.log" 2>&1 \
     && git -C "${dir}" reset --hard "origin/${branch}" >>"${LOG_DIR}/panel-update.log" 2>&1; then
    log "${label} up to date: $(git -C "${dir}" rev-parse --short HEAD 2>/dev/null)"
  else
    log "WARN: ${label} update check failed — keeping the existing copy (see panel-update.log)"
  fi
}
PANEL_DIR="${CN_VOL}/comfyui-mcp-panel"
if [ "${PANEL_AUTO_UPDATE}" = "1" ]; then
  ff_baked_node "comfyui-mcp-panel" "Agent Panel"
  ff_baked_node "ComfyUI-Crystools" "Crystools"
else
  log "baked-node auto-update disabled (PANEL_AUTO_UPDATE=${PANEL_AUTO_UPDATE})"
fi

# (b.2) PANEL INTEGRITY CHECK + SELF-HEAL. A past ENOSPC (or any interrupted
#     copy) can leave the panel on the volume as a full file tree of 0-BYTE
#     files — ComfyUI then lists the node but the panel tab never loads, and
#     because the volume PERSISTS, every later redeploy looks just as broken.
#     Verify the three load-bearing files; on failure re-clone from scratch.
panel_ok() {  # $1 = panel dir
  [ -s "$1/__init__.py" ] && [ -s "$1/pyproject.toml" ] \
    && [ -s "$1/web/js/comfyui-mcp-panel.js" ]
}
if ! panel_ok "${PANEL_DIR}"; then
  log "Agent Panel on the volume is BROKEN (missing/0-byte files) — self-healing…"
  rm -rf "${PANEL_DIR}"
  if git clone --depth 1 ${PANEL_REF:+--branch "${PANEL_REF}"} "${PANEL_REPO}" "${PANEL_DIR}" \
       >>"${LOG_DIR}/panel-update.log" 2>&1 && panel_ok "${PANEL_DIR}"; then
    log "Agent Panel re-cloned OK: $(git -C "${PANEL_DIR}" rev-parse --short HEAD 2>/dev/null)"
  elif [ -d "${CN_SEED}/comfyui-mcp-panel" ] \
       && rm -rf "${PANEL_DIR}" \
       && cp -rf "${CN_SEED}/comfyui-mcp-panel" "${PANEL_DIR}" 2>>"${LOG_DIR}/panel-update.log" \
       && panel_ok "${PANEL_DIR}"; then
    log "Agent Panel restored from the image seed (offline fallback)."
  else
    log "============================================================================"
    log "ERROR: could not repair the Agent Panel on ${CN_VOL}. Most common cause: the"
    log "  network volume is FULL ($(free_mb "${WORKSPACE}") MB free) — writes create"
    log "  0-byte files. Free/grow the volume and restart the pod; the panel will"
    log "  self-heal on the next boot. Details: ${LOG_DIR}/panel-update.log"
    log "============================================================================"
  fi
fi

# Zero-byte sweep across ALL nodes on the volume — the same ENOSPC event that
# empties the panel empties user-installed nodes too. We can't safely re-clone
# arbitrary nodes (unknown sources), but we CAN say exactly which are broken.
BROKEN_NODES="$(find "${CN_VOL}" -mindepth 2 -maxdepth 2 -name '__init__.py' -size 0 2>/dev/null \
  | sed 's|.*/custom_nodes/||; s|/__init__.py$||' | sort -u | tr '\n' ' ')"
if [ -n "${BROKEN_NODES// /}" ]; then
  log "WARN: custom nodes with 0-byte __init__.py (broken; reinstall via Manager): ${BROKEN_NODES}"
fi

# (c) Reinstall custom-node Python deps into the venv from the persistent cache.
#     Runs every boot (the venv is ephemeral); the cache makes it fast. Best-effort
#     per node — a broken requirements.txt must not take the whole pod down.
VPIP="${COMFY_HOME}/venv/bin/pip"
if [ -x "${VPIP}" ]; then
  cn_count=0
  for req in "${CN_VOL}"/*/requirements.txt; do
    [ -f "${req}" ] || continue          # literal glob when no matches → skip
    node_name="$(basename "$(dirname "${req}")")"
    log "custom-node deps: installing ${node_name} requirements…"
    if "${VPIP}" install --no-input --disable-pip-version-check -r "${req}" \
         >>"${LOG_DIR}/custom-node-deps.log" 2>&1; then
      cn_count=$((cn_count + 1))
    else
      log "WARN: deps for ${node_name} failed (see custom-node-deps.log; node may not load)"
    fi
  done
  log "custom-node deps: processed ${cn_count} node requirement set(s) (cache: ${PIP_CACHE_DIR})"
else
  log "WARN: venv pip not found at ${VPIP} — skipping custom-node dep install"
fi

# -----------------------------------------------------------------------------
# 5. Launch ComfyUI from the BAKED venv (image), pointed at the volume dirs.
#    Invoke the venv python by ABSOLUTE PATH (no `activate` needed).
#    Per-directory flags keep user/input/output on /workspace. --models-directory
#    makes ComfyUI's folder_paths.models_dir the persistent volume root, which is
#    required for Manager's explicit relative save_path values; the extra path map
#    still puts every category's volume subfolder first. custom_nodes are symlinked
#    onto the volume in §4.5 above (so runtime installs persist); we do NOT use
#    --base-directory (it would relocate the whole tree, incl. the venv).
# -----------------------------------------------------------------------------
VPY="${COMFY_HOME}/venv/bin/python"
if [ ! -x "${VPY}" ]; then
  log "FATAL: baked venv python not found at ${VPY} — image build problem."
  log "       Holding pod open for debug (the image software did not bake)."
  exec sleep infinity
fi

COMFY_LOG="${LOG_DIR}/comfyui.log"

# ComfyUI 0.27's sqlite DB defaults to <base>/user/comfyui.db (= ${COMFY_HOME}/user),
# computed from the BASE dir — NOT --user-directory — and that dir isn't in the
# image, so init fails ("unable to open database file"). Create it so the DB
# initializes cleanly (local, ephemeral index; re-created per boot).
mkdir -p "${COMFY_HOME}/user"

# --enable-cors-header is REQUIRED for RunPod-proxy access. ComfyUI's
# origin_only_middleware (server.py) returns 403 for any request whose
# `Sec-Fetch-Site: cross-site` — exactly what a browser sends when it reaches
# ComfyUI THROUGH the RunPod proxy (cross-origin to the proxy domain). That 403s
# the whole UI ("won't render") even though curl (no Sec-Fetch headers) works.
# --enable-cors-header swaps that middleware for the CORS one, letting the
# proxied browser through.
# PERF: attention backend PROBED, not assumed — the cu130 perf variant bakes
# SageAttention 2.2 (+ triton backend) but the default cu128 build ships neither
# (no linux cu128 wheels exist for this torch line). Passing --use-sage-attention
# without the package would crash ComfyUI at startup, so import-check the venv
# and pick the matching flags. Override either way via COMFY_EXTRA_ARGS.
if "${COMFY_HOME}/venv/bin/python" -c "import sageattention" >/dev/null 2>&1; then
  ATTN_ARGS=(--use-sage-attention --enable-triton-backend)
  log "SageAttention baked in this image — launching with --use-sage-attention"
else
  ATTN_ARGS=(--use-pytorch-cross-attention)
  log "no SageAttention in this image (cu128 broad-compat build) — using PyTorch cross-attention"
fi
ARGS=(--listen 0.0.0.0 --port "${COMFY_PORT}"
      --enable-manager --enable-cors-header
      "${ATTN_ARGS[@]}"
      --user-directory  "${USER_DIR}"
      --models-directory "${MODELS_DIR}"
      --input-directory "${INPUT_DIR}"
      --output-directory "${OUTPUT_DIR}")
# Load the volume model map only if the file exists (it's baked, but be defensive).
[ -f "${EXTRA_MODEL_PATHS}" ] && ARGS+=(--extra-model-paths-config "${EXTRA_MODEL_PATHS}")
# shellcheck disable=SC2206
[ -n "${COMFY_EXTRA_ARGS}" ] && ARGS+=(${COMFY_EXTRA_ARGS})

# ---- The one line a user needs to drive this pod's Agent Panel ---------------
#
# The panel already builds this command itself (its connectCommand() appends the
# page's own origin whenever the page is served over https), but only INSIDE the
# sidebar's setup card. A user who has not opened that card - or who is looking at
# the "starting" page because ComfyUI is still booting - never sees it, and the
# panel's Bridge URL field reads ws://127.0.0.1:9199, which looks like the pod is
# misconfigured. It is not: the orchestrator runs on the USER's machine and tunnels
# back to the panel, so this proxy URL is the only part they cannot guess.
POD_URL=""
[ -n "${RUNPOD_POD_ID:-}" ] && POD_URL="https://${RUNPOD_POD_ID}-3000.proxy.runpod.net"
POD_URL_LOG="${POD_URL:-https://<pod-id>-3000.proxy.runpod.net}"

# Stamp it into the "ComfyUI is starting…" page (nginx serves that whenever :3001
# is not reachable yet) - the FIRST thing anyone sees on a cold pod. Best-effort:
# a missing or already-substituted page must never keep the pod from booting.
STARTING_PAGE="${STARTING_PAGE:-/usr/share/nginx/html/readme.html}"
if [ -f "${STARTING_PAGE}" ]; then
  if sed -i "s|__POD_URL__|${POD_URL:-https://&lt;pod-id&gt;-3000.proxy.runpod.net}|g" "${STARTING_PAGE}" 2>/dev/null; then
    log "starting page stamped with the connect command (${STARTING_PAGE})"
  else
    log "WARN: could not stamp ${STARTING_PAGE} — the starting page keeps its placeholder"
  fi
fi

cd "${COMFY_HOME}"
log "launching ComfyUI: ${VPY} main.py ${ARGS[*]}"
log "  software   : ${COMFY_HOME}        (image; immutable, fast local import)"
log "  user dir   : ${USER_DIR}          (volume; workflows + settings)"
log "  models     : ${MODELS_DIR}        (volume; downloads persist here)"
log "  input/out  : ${INPUT_DIR} / ${OUTPUT_DIR}  (volume)"
log "  HTTP (nginx): :3000  ->  ComfyUI :${COMFY_PORT}"
log "  RunPod proxy: ${POD_URL_LOG}"
nohup "${VPY}" main.py "${ARGS[@]}" >>"${COMFY_LOG}" 2>&1 &
COMFY_PID=$!
log "ComfyUI started (pid=${COMFY_PID}); streaming ${COMFY_LOG}"

# Printed AFTER the launch line so it is the last thing in the boot log before the
# ComfyUI stream takes over - i.e. what a user actually sees in the RunPod console.
log ""
log "──────────────────────────────────────────────────────────────────────"
log " Agent Panel — run this on YOUR machine to connect it to this pod:"
log ""
log "   npx -y comfyui-mcp@latest connect ${POD_URL_LOG}"
log ""
log " The agent runs on your own Claude / ChatGPT / Gemini login. Nothing"
log " extra is installed here. Then click Connect in the Agent sidebar."
log "──────────────────────────────────────────────────────────────────────"
log ""

# Stream ComfyUI's log to the pod console and hold the pod open. If tail is ever
# killed, the base's `sleep infinity` still keeps the pod alive.
exec tail -n +1 -F "${COMFY_LOG}"