#!/bin/bash

# Vast.ai ComfyUI Provisioning Script
# Updated: Includes update command and model verification

echo "Starting provisioning script..."

# 1. Ensure all required directories exist
mkdir -p /workspace/ComfyUI/models/{text_encoders,checkpoints,vae,clip,loras,diffusion_models,upscale_models}
mkdir -p /workspace/ComfyUI/models/ultralytics/bbox

# 2. Define Helper for Civitai Downloads
AUTH_HEADER="Authorization: Bearer $CIVITAI_TOKEN"

# 3. === MODEL DOWNLOADS (Hugging Face) ===
echo "Checking Hugging Face models..."

# Flux & Qwen Base Models
[ ! -f /workspace/ComfyUI/models/text_encoders/t5xxl_fp8_e4m3fn_scaled.safetensors ] && wget -P /workspace/ComfyUI/models/text_encoders https://huggingface.co/comfyanonymous/flux_text_encoders/resolve/main/t5xxl_fp8_e4m3fn_scaled.safetensors
[ ! -f /workspace/ComfyUI/models/checkpoints/Qwen-Rapid-AIO-SFW-v23.safetensors ] && wget -P /workspace/ComfyUI/models/checkpoints https://huggingface.co/Phr00t/Qwen-Image-Edit-Rapid-AIO/resolve/main/v23/Qwen-Rapid-AIO-SFW-v23.safetensors
[ ! -f /workspace/ComfyUI/models/checkpoints/Qwen-Rapid-AIO-NSFW-v23.safetensors ] && wget -P /workspace/ComfyUI/models/checkpoints https://huggingface.co/Phr00t/Qwen-Image-Edit-Rapid-AIO/resolve/main/v23/Qwen-Rapid-AIO-NSFW-v23.safetensors
[ ! -f /workspace/ComfyUI/models/vae/qwen_image_vae.safetensors ] && wget -P /workspace/ComfyUI/models/vae "https://huggingface.co/Comfy-Org/Qwen-Image_ComfyUI/resolve/main/split_files/vae/qwen_image_vae.safetensors"
[ ! -f /workspace/ComfyUI/models/clip/qwen_3_4b.safetensors ] && wget -P /workspace/ComfyUI/models/clip https://huggingface.co/Comfy-Org/z_image_turbo/resolve/main/split_files/text_encoders/qwen_3_4b.safetensors
[ ! -f /workspace/ComfyUI/models/text_encoders/qwen_2.5_vl_7b_fp8_scaled.safetensors ] && wget -P /workspace/ComfyUI/models/text_encoders "https://huggingface.co/Comfy-Org/Qwen-Image_ComfyUI/resolve/main/split_files/text_encoders/qwen_2.5_vl_7b_fp8_scaled.safetensors"
[ ! -f /workspace/ComfyUI/models/diffusion_models/z_image_turbo_bf16.safetensors ] && wget -P /workspace/ComfyUI/models/diffusion_models https://huggingface.co/Comfy-Org/z_image_turbo/resolve/main/split_files/diffusion_models/z_image_turbo_bf16.safetensors
[ ! -f /workspace/ComfyUI/models/vae/flux_vae.safetensors ] && wget -P /workspace/ComfyUI/models/vae https://huggingface.co/StableDiffusionVN/Flux/resolve/main/Vae/flux_vae.safetensors
[ ! -f /workspace/ComfyUI/models/ultralytics/bbox/FacesV1.pt ] && wget -P /workspace/ComfyUI/models/ultralytics/bbox "https://huggingface.co/ashllay/YOLO_Models/resolve/e07b01219ff1807e1885015f439d788b038f49bd/bbox/FacesV1.pt"
[ ! -f /workspace/ComfyUI/models/upscale_models/1x-ITF-SkinDiffDetail-Lite-v1.pth ] && wget -P /workspace/ComfyUI/models/upscale_models/ "https://huggingface.co/cozyga/test/resolve/main/1x-ITF-SkinDiffDetail-Lite-v1.pth"

# LoRAs (Hugging Face)
LORA_DIR="/workspace/ComfyUI/models/loras"
[ ! -f "$LORA_DIR/Z-Isla_3K_0002_000003500.safetensors" ] && wget -P "$LORA_DIR" https://huggingface.co/cozyga/test/resolve/main/Z-Isla_3K_0002_000003500.safetensors
[ ! -f "$LORA_DIR/z-kyla-v1.safetensors" ] && wget -P "$LORA_DIR" https://huggingface.co/cozyga/test/resolve/main/z-kyla-v1.safetensors
[ ! -f "$LORA_DIR/m1a_zit_3.safetensors" ] && wget -P "$LORA_DIR" https://huggingface.co/cozyga/test/resolve/main/m1a_zit_3.safetensors
[ ! -f "$LORA_DIR/c3lia_v1_000003600.safetensors" ] && wget -P "$LORA_DIR" https://huggingface.co/cozyga/test/resolve/main/c3lia_v1_000003600.safetensors
[ ! -f "$LORA_DIR/c3lia_v1_000003900.safetensors" ] && wget -P "$LORA_DIR" https://huggingface.co/cozyga/test/resolve/main/c3lia_v1_000003900.safetensors
[ ! -f "$LORA_DIR/qwen_c3lia_000003600.safetensors" ] && wget -P "$LORA_DIR" https://huggingface.co/cozyga/test/resolve/main/qwen_c3lia_000003600.safetensors
[ ! -f "$LORA_DIR/l0ki_zit_v1_000003000.safetensors" ] && wget -P "$LORA_DIR" https://huggingface.co/XLB666/qwenzitloras/resolve/main/l0ki_zit_v1_000003000.safetensors
[ ! -f "$LORA_DIR/l0ki_qwen_v1_000002800.safetensors" ] && wget -P "$LORA_DIR" https://huggingface.co/XLB666/qwenzitloras/resolve/main/l0ki_qwen_v1_000002800.safetensors
[ ! -f "$LORA_DIR/s0fia_v2_000004800.safetensors" ] && wget -P "$LORA_DIR" https://huggingface.co/cozyga/test/resolve/main/s0fia_v2_000004800.safetensors
[ ! -f "$LORA_DIR/s0fia_v2_000003700.safetensors" ] && wget -P "$LORA_DIR" https://huggingface.co/cozyga/test/resolve/main/s0fia_v2_000003700.safetensors
[ ! -f "$LORA_DIR/s0fia_v2_000004300.safetensors" ] && wget -P "$LORA_DIR" https://huggingface.co/cozyga/test/resolve/main/s0fia_v2_000004300.safetensors

# 4. === MODEL DOWNLOADS (Civitai) ===
echo "Checking Civitai models..."

[ ! -f /workspace/ComfyUI/models/loras/zittitslider.safetensors ] && \
    curl -L -H "$AUTH_HEADER" -o /workspace/ComfyUI/models/loras/zittitslider.safetensors "https://civitai.com/api/download/models/2478366"
[ ! -f /workspace/ComfyUI/models/loras/QwenEditVersatilePoses.safetensors ] && \
    curl -L -H "$AUTH_HEADER" -o /workspace/ComfyUI/models/loras/QwenEditVersatilePoses.safetensors "https://civitai.com/api/download/models/2457989?type=Model&format=SafeTensor"
[ ! -f /workspace/ComfyUI/models/loras/sdxlcloseuppussy.safetensors ] && \
    curl -L -H "$AUTH_HEADER" -o /workspace/ComfyUI/models/loras/sdxlcloseuppussy.safetensors "https://civitai.com/api/download/models/2376235?type=Model&format=SafeTensor"
[ ! -f /workspace/ComfyUI/models/vae/zImage_vae.safetensors ] && \
    curl -L -H "$AUTH_HEADER" -o /workspace/ComfyUI/models/vae/zImage_vae.safetensors "https://civitai.com/api/download/models/2442479"
[ ! -f /workspace/ComfyUI/models/loras/amateurphotography.safetensors ] && \
    curl -L -H "$AUTH_HEADER" -o /workspace/ComfyUI/models/loras/amateurphotography.safetensors "https://civitai.com/api/download/models/2363467"
[ ! -f /workspace/ComfyUI/models/loras/nextscene.safetensors ] && \
    curl -L -H "$AUTH_HEADER" -o /workspace/ComfyUI/models/loras/nextscene.safetensors "https://civitai.com/api/download/models/2337214?type=Model&format=SafeTensor"

# 5. === UPDATE AND START ===
echo "Provisioning complete. Updating ComfyUI..."

cd /workspace/ComfyUI || exit
git pull

echo "ComfyUI is updated. Starting now..."
# Ensure the script ends by handing over to the ComfyUI process
python main.py --listen