#!/usr/bin/env bash
# 팟에서 실행한다. tools/pod_bootstrap.py script 가 만든다.
set -euo pipefail

COMFY_ROOT="${COMFY_ROOT:-}"
STAGE="$(mktemp -d)"
trap 'rm -rf "$STAGE"' EXIT

# ComfyUI 본체 경로는 이미지마다 다르다(runpod/comfyui 는 /workspace/runpod-slim/ComfyUI).
if [ -z "$COMFY_ROOT" ]; then
    for candidate in /workspace/ComfyUI /workspace/runpod-slim/ComfyUI /ComfyUI /root/ComfyUI /app/ComfyUI; do
        if [ -f "$candidate/main.py" ]; then COMFY_ROOT="$candidate"; break; fi
    done
fi
if [ ! -f "$COMFY_ROOT/main.py" ]; then
    echo "ComfyUI 본체를 찾지 못했습니다. COMFY_ROOT=/실제/경로 로 지정해 다시 실행하세요." >&2
    exit 1
fi
echo "ComfyUI: $COMFY_ROOT"

# ComfyUI 가 쓰는 인터프리터(이미지에 전용 venv 가 있다).
# glob 을 ls 로 풀면 매치가 없을 때 exit 2 가 되고, pipefail 때문에 조용히 죽는다.
PY=""
for candidate in "$COMFY_ROOT"/.venv*/bin/python; do
    if [ -x "$candidate" ]; then PY="$candidate"; break; fi
done
if [ -z "$PY" ]; then PY="$(command -v python3 || command -v python || true)"; fi
# hf CLI 는 이미지에 없을 수 있다. 없으면 같은 인터프리터의 huggingface_hub 로 받는다.
HF="${HF:-}"
for candidate in "$COMFY_ROOT"/.venv*/bin/hf; do
    if [ -x "$candidate" ]; then HF="$candidate"; break; fi
done
if [ -z "$HF" ]; then HF="$(command -v hf || true)"; fi
if [ -z "$HF" ]; then echo "hf CLI 가 없어 파이썬 API로 내려받습니다." >&2; fi

fetch() {
    local repo="$1" name="$2" dest="$3"
    if [ -n "$HF" ]; then
        "$HF" download "$repo" "$name" --local-dir "$dest" >/dev/null
    else
        "$PY" -c 'import sys;from huggingface_hub import hf_hub_download;hf_hub_download(sys.argv[1],sys.argv[2],local_dir=sys.argv[3])' "$repo" "$name" "$dest" >/dev/null
    fi
}

# 모델: 공개 저장소라 토큰이 필요 없다. 받은 뒤 이름으로 찾아 제자리에 옮긴다.
pull() {
    local repo="$1" name="$2" target="$3" src
    fetch "$repo" "$name" "$STAGE/$repo"
    src="$(find "$STAGE/$repo" -name "$name" -print -quit || true)"
    if [ -z "$src" ]; then echo "HF에 없음: $repo/$name" >&2; return 1; fi
    install -D -m 644 "$src" "$COMFY_ROOT/models/$target"
    # 빈 파일이면 실패로 끝낸다 - 파일명만 보는 점검은 이걸 못 잡는다(실제로 못 잡았다).
    if [ ! -s "$COMFY_ROOT/models/$target" ]; then
        echo "내려받기 실패(빈 파일): $target" >&2; return 1
    fi
    # 하위 폴더 모델은 역슬래시 이름 심볼릭 링크도 만든다. ComfyUI 는 폴더 목록을
    # os.sep 로 만들기 때문에(folder_paths.py) 윈도우에서 저장된 워크플로의
    # 'bodyparts\x' 값이 리눅스 팟에서는 어떤 콤보와도 안 맞는다. 링크를 만들어
    # 원본 워크플로를 고치지 않고도 양쪽에서 같은 값이 매칭되게 한다.
    # 하위 폴더가 있는 경우에만 별칭을 만든다(*/* 로 두면 unet/x.gguf 처럼
    # 폴더 바로 아래 파일도 매치되어 제자리를 가리키는 링크로 실제 파일을 지운다).
    case "$target" in
        */*/*)
            local root="${target%%/*}" rest="${target#*/}" win
            win="$(printf %s "$rest" | tr "/" "\\")"
            ln -sf "$rest" "$COMFY_ROOT/models/$root/$win"
            ;;
    esac
}

mkdir -p "$COMFY_ROOT/models/unet" "$COMFY_ROOT/models/vae" \
         "$COMFY_ROOT/models/text_encoders" "$COMFY_ROOT/models/loras" \
         "$COMFY_ROOT/custom_nodes"

pull bjooo/wan2.2_loras "SVI_v2_PRO_Wan2.2-I2V-A14B_HIGH_lora_rank_128_fp16.safetensors" "loras/SVI_v2_PRO_Wan2.2-I2V-A14B_HIGH_lora_rank_128_fp16.safetensors"
pull bjooo/wan2.2_loras "SVI_v2_PRO_Wan2.2-I2V-A14B_LOW_lora_rank_128_fp16.safetensors" "loras/SVI_v2_PRO_Wan2.2-I2V-A14B_LOW_lora_rank_128_fp16.safetensors"
pull bjooo/wan2.2_loras "Wan2.1_I2V_14B_FusionX_LoRA.safetensors" "loras/Wan2.1_I2V_14B_FusionX_LoRA.safetensors"
pull bjooo/wan2.2_loras "Wan_2_2_I2V_A14B_HIGH_lightx2v_4step_lora_v1030_rank_64_bf16.safetensors" "loras/Wan_2_2_I2V_A14B_HIGH_lightx2v_4step_lora_v1030_rank_64_bf16.safetensors"
pull bjooo/wan2.2_loras "Wan_2_2_I2V_A14B_LOW_lightx2v_4step_lora_260412_rank_256_fp16.safetensors" "loras/Wan_2_2_I2V_A14B_LOW_lightx2v_4step_lora_260412_rank_256_fp16.safetensors"
pull bjooo/wan2.2_loras "2.2-I2V Reveal Penis_000003000_high_noise.safetensors" "loras/bodyparts/2.2-I2V Reveal Penis_000003000_high_noise.safetensors"
pull bjooo/wan2.2_loras "2.2-I2V Reveal Penis_000003000_low_noise.safetensors" "loras/bodyparts/2.2-I2V Reveal Penis_000003000_low_noise.safetensors"
pull bjooo/wan2.2_loras "PENISLORA_22_i2v_HIGH_e320.safetensors" "loras/bodyparts/PENISLORA_22_i2v_HIGH_e320.safetensors"
pull bjooo/wan2.2_loras "PENISLORA_22_i2v_LOW_e496.safetensors" "loras/bodyparts/PENISLORA_22_i2v_LOW_e496.safetensors"
pull bjooo/wan2.2_loras "erect_penis_epoch_80.safetensors" "loras/bodyparts/erect_penis_epoch_80.safetensors"
pull bjooo/wan2.2_loras "masturbation_cumshot_v1.1_e310.safetensors" "loras/masturbation/masturbation_cumshot_v1.1_e310.safetensors"
pull bjooo/clips "nsfw_wan_umt5-xxl_fp8_scaled.safetensors" "text_encoders/nsfw_wan_umt5-xxl_fp8_scaled.safetensors"
pull bjooo/unet "Wan2.2-I2V-A14B-HighNoise-Q8_0.gguf" "unet/Wan2.2-I2V-A14B-HighNoise-Q8_0.gguf"
pull bjooo/unet "Wan2.2-I2V-A14B-LowNoise-Q8_0.gguf" "unet/Wan2.2-I2V-A14B-LowNoise-Q8_0.gguf"
pull bjooo/clips "Wan2_1_VAE_bf16.safetensors" "vae/Wan2_1_VAE_bf16.safetensors"

# 이미지에 같은 계열 팩이 들어 있으면 custom_nodes 밖으로 옮긴다. 리눅스는
# 대소문자를 구분하므로 ComfyUI-KJNodes 와 우리 comfyui-kjnodes 가 공존하면
# 같은 클래스가 두 번 등록된다. (.disabled 접미사는 ComfyUI 가 무시한다는 보장이 약하다)
mkdir -p "$COMFY_ROOT/.parked_nodes"
for collide in ComfyUI-KJNodes; do
    if [ -d "$COMFY_ROOT/custom_nodes/$collide" ]; then
        mv "$COMFY_ROOT/custom_nodes/$collide" "$COMFY_ROOT/.parked_nodes/$collide"
        echo "충돌 팩을 옮겼습니다: $collide -> .parked_nodes/"
    fi
done

# 노드: 로컬 설치본 번들 + 클론 가능한 팩
fetch bjooo/customNode node_bundle.tar.gz "$STAGE/bundle"
tar -xzf "$STAGE/bundle/node_bundle.tar.gz" -C "$COMFY_ROOT/custom_nodes"

# 키를 노드 폴더 .env 로도 남긴다. 팟 컨테이너가 재시작되면(이미지가 띄우므로)
# 실행 환경변수는 초기화되지만 .env 는 볼륨에 남고 prestartup_script.py 가 읽는다.
if [ -n "${OPENROUTER_API_KEY:-}" ] && [ -d "$COMFY_ROOT/custom_nodes/dolphin_nodes" ]; then
    printf "OPENROUTER_API_KEY=%s\n" "$OPENROUTER_API_KEY" > "$COMFY_ROOT/custom_nodes/dolphin_nodes/.env"
    chmod 600 "$COMFY_ROOT/custom_nodes/dolphin_nodes/.env"
    echo "OPENROUTER_API_KEY 를 노드 폴더 .env 에 남겼습니다(재시작에도 유지)."
fi
[ -d "$COMFY_ROOT/custom_nodes/RES4LYF" ] || git clone --depth 1 https://github.com/ClownsharkBatwing/RES4LYF "$COMFY_ROOT/custom_nodes/RES4LYF"

# 팩별 파이썬 의존성은 ComfyUI 가 쓰는 인터프리터에 넣어야 한다.
for req in "$COMFY_ROOT"/custom_nodes/*/requirements.txt; do
    [ -f "$req" ] || continue
    "$PY" -m pip install -q -r "$req" || echo "의존성 설치 실패: $req" >&2
done

# 감독 노드가 실행 중 OpenRouter를 부른다. 없으면 렌더가 시작하자마자 죽는다.
if [ -z "${OPENROUTER_API_KEY:-}" ]; then
    echo "경고: OPENROUTER_API_KEY 가 없습니다. SVI 감독 노드가 즉시 실패합니다." >&2
fi

# 새 노드를 등록하려면 ComfyUI 프로세스만 다시 띄운다. 팟 재시작은 컨테이너 디스크를
# 통째로 초기화해서 방금 받은 모델까지 다 날아간다(2026-09-30 실측). 절대 쓰지 말 것.
ARGS="--listen 0.0.0.0 --port 8188 --enable-cors-header"
if command -v pkill >/dev/null 2>&1; then
    pkill -f "main.py --listen" || true
    sleep 3
else
    echo "pkill 이 없어 자동 재시작을 건너뜁니다. ComfyUI를 직접 재시작하세요." >&2
fi
cd "$COMFY_ROOT"
nohup "$PY" main.py $ARGS > /tmp/comfyui_bootstrap.log 2>&1 &
for _ in $(seq 1 60); do
    sleep 2
    if curl -fsS -o /dev/null "http://127.0.0.1:8188/system_stats"; then
        echo "완료. ComfyUI가 새 노드로 올라왔습니다. 사이드바 Runpod 렌더 탭에서 프리플라이트를 돌리세요."
        exit 0
    fi
done
echo "ComfyUI가 2분 안에 응답하지 않았습니다. /tmp/comfyui_bootstrap.log 를 확인하세요." >&2
tail -20 /tmp/comfyui_bootstrap.log >&2
exit 1