#!/usr/bin/env bash
# Container CMD for the serverless image: boot ComfyUI headless, then run the
# RunPod handler. Models + code are baked; nothing is fetched at cold start
# unless SERVERLESS_REFRESH=true (then handler/workflows/params are re-pulled
# from the HF repo best-effort, for zero-rebuild iteration).
set -u

COMFY_DIR="${COMFY_DIR:-/workspace/ComfyUI}"
COMFY_PORT="${COMFY_PORT:-7865}"
SRC="${SERVERLESS_SRC:-/opt/serverless/src}"
export PYTHONPATH="/opt/serverless/lib${PYTHONPATH:+:$PYTHONPATH}"

if [ "${SERVERLESS_REFRESH:-false}" = "true" ]; then
    echo "[start] refreshing serverless code from HF repo (SERVERLESS_REFRESH=true)"
    tok="${HUGGING_FACE_ACCESS_TOKEN:-${HF_TOKEN:-}}"
    auth=(); [ -n "$tok" ] && auth=(-H "Authorization: Bearer $tok")
    base="https://huggingface.co/aleph65/ComfyUI/resolve/main/serverless"
    while read -r rel; do
        tmp=$(mktemp)
        if curl -fsSL --connect-timeout 5 --max-time 20 "${auth[@]}" \
                "$base/$rel" -o "$tmp" 2>/dev/null; then
            mkdir -p "$SRC/$(dirname "$rel")" && mv "$tmp" "$SRC/$rel"
            echo "[start]   refreshed $rel"
        else
            rm -f "$tmp"
        fi
    done < <(cd "$SRC" && find . -name '*.py' -o -name '*.json' | sed 's|^\./||')
fi

cd "$COMFY_DIR"
# shellcheck disable=SC2086 — COMFY_ARGS is intentionally word-split
python -u main.py --listen 127.0.0.1 --port "$COMFY_PORT" \
    --disable-auto-launch ${COMFY_ARGS:-} \
    > /comfyui.log 2>&1 &

exec python -u "$SRC/handler.py"