#!/usr/bin/env bash
set -uo pipefail

# ── ComfyUI vast.ai Setup Script ──────────────────────────────
# Run from inside the cloned ComfyUI directory.
# Usage: git clone <repo> && cd ComfyUI && bash setup-vastai.sh
#
# Expects: NVIDIA GPU, Python 3.10+, git, wget

SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
MODELS_DIR="$SCRIPT_DIR/models"
FAILED_DOWNLOADS=()
FAILED_NODES=()
CURRENT_DOWNLOAD=""

cleanup() {
    echo ""
    if [ -n "$CURRENT_DOWNLOAD" ] && [ -f "$CURRENT_DOWNLOAD" ]; then
        rm -f "$CURRENT_DOWNLOAD"
        err "Cleaned up partial download: $(basename "$CURRENT_DOWNLOAD")"
    fi
    err "Setup interrupted. Re-run the script to resume — completed downloads will be skipped."
    exit 130
}
trap cleanup INT TERM

# ── Helpers ───────────────────────────────────────────────────
log()  { echo -e "\033[1;32m[setup]\033[0m $*"; }
warn() { echo -e "\033[1;33m[warn]\033[0m  $*" >&2; }
err()  { echo -e "\033[1;31m[error]\033[0m $*" >&2; }
die()  { err "$*"; exit 1; }

dl() {
    local url="$1" dest="$2" name
    name="$(basename "$dest")"
    if [ -f "$dest" ]; then
        local size
        size=$(stat -c%s "$dest" 2>/dev/null || stat -f%z "$dest" 2>/dev/null || echo 0)
        if [ "$size" -gt 0 ]; then
            log "  Skip (exists): $name"
            return 0
        else
            warn "  Removing empty file: $name"
            rm -f "$dest"
        fi
    fi
    mkdir -p "$(dirname "$dest")"
    log "  Downloading: $name"
    local tmp="${dest}.part"
    CURRENT_DOWNLOAD="$tmp"
    if wget -q --show-progress -O "$tmp" "$url"; then
        mv "$tmp" "$dest"
        CURRENT_DOWNLOAD=""
    else
        rm -f "$tmp"
        CURRENT_DOWNLOAD=""
        err "  Failed to download: $name"
        FAILED_DOWNLOADS+=("$name ($url)")
        return 1
    fi
}

symlink() {
    local target="$1" link="$2"
    mkdir -p "$(dirname "$link")"
    if [ ! -f "$target" ]; then
        warn "  Symlink target missing: $(basename "$target") — skipping"
        return 1
    fi
    ln -sf "$target" "$link"
}

# ── Preflight Checks ────────────────────────────────────────
cd "$SCRIPT_DIR"

[ -f "main.py" ] || die "Not in a ComfyUI directory (main.py not found). Run from inside the repo."

for cmd in git python3 wget; do
    command -v "$cmd" &>/dev/null || die "'$cmd' is required but not found."
done

log "Starting ComfyUI setup in: $SCRIPT_DIR"

# ── Git LFS ──────────────────────────────────────────────────
if git lfs version &>/dev/null; then
    log "Pulling LFS files (Civitai loras)..."
    git lfs install > /dev/null 2>&1
    git lfs pull || warn "LFS pull failed — Civitai loras may be missing"
else
    log "git-lfs not found, installing..."
    if apt-get update -qq && apt-get install -y -qq git-lfs > /dev/null 2>&1; then
        git lfs install > /dev/null 2>&1
        git lfs pull || warn "LFS pull failed — Civitai loras may be missing"
    else
        warn "Could not install git-lfs — Civitai loras will not be available"
    fi
fi

# ── Python Environment ───────────────────────────────────────
log "Setting up Python environment..."
if [ ! -d "$SCRIPT_DIR/venv" ]; then
    python3 -m venv venv || { err "Failed to create virtual environment"; }
fi
# shellcheck disable=SC1091
source venv/bin/activate

# Install PyTorch with CUDA
log "Installing PyTorch with CUDA..."
pip install --quiet --upgrade torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu124 \
    || err "Failed to install PyTorch — ComfyUI will not work without it"

# Verify CUDA is available
if python3 -c "import torch; assert torch.cuda.is_available(), 'no cuda'" 2>/dev/null; then
    log "PyTorch CUDA: OK ($(python3 -c 'import torch; print(torch.cuda.get_device_name(0))'))"
else
    warn "PyTorch installed but CUDA not available — generation will be very slow"
fi

# Install ComfyUI requirements
log "Installing ComfyUI requirements..."
pip install --quiet --upgrade -r requirements.txt \
    || err "Failed to install some ComfyUI requirements"

# ── Custom Node Requirements ─────────────────────────────────
log "Installing custom node requirements..."
for req in custom_nodes/*/requirements.txt; do
    [ -f "$req" ] || continue
    node_name="$(basename "$(dirname "$req")")"
    log "  $node_name"
    if ! pip install --quiet --upgrade -r "$req" 2>/dev/null; then
        warn "  Some deps failed for $node_name"
        FAILED_NODES+=("$node_name")
    fi
done

# ── Download Models ───────────────────────────────────────────
log "Downloading models from HuggingFace..."
log ""

# Checkpoints
log "[ Checkpoints ]"
dl "https://huggingface.co/Comfy-Org/flux1-dev/resolve/main/flux1-dev-fp8.safetensors" \
   "$MODELS_DIR/checkpoints/flux1-dev-fp8.safetensors"

# CLIP / Text Encoders
log "[ CLIP & Text Encoders ]"
dl "https://huggingface.co/comfyanonymous/flux_text_encoders/resolve/main/clip_l.safetensors" \
   "$MODELS_DIR/clip/clip_l.safetensors"
dl "https://huggingface.co/comfyanonymous/flux_text_encoders/resolve/main/t5xxl_fp8_e4m3fn.safetensors" \
   "$MODELS_DIR/clip/t5xxl_fp8_e4m3fn.safetensors"
dl "https://huggingface.co/QuanSun/EVA-CLIP/resolve/main/EVA02_CLIP_L_336_psz14_s6B.pt" \
   "$MODELS_DIR/clip/EVA02_CLIP_L_336_psz14_s6B.pt"

# Symlink text_encoders -> clip (same files)
mkdir -p "$MODELS_DIR/text_encoders"
symlink "$MODELS_DIR/clip/clip_l.safetensors" "$MODELS_DIR/text_encoders/clip_l.safetensors"
symlink "$MODELS_DIR/clip/t5xxl_fp8_e4m3fn.safetensors" "$MODELS_DIR/text_encoders/t5xxl_fp8_e4m3fn.safetensors"

# CLIP Vision
log "[ CLIP Vision ]"
dl "https://huggingface.co/Comfy-Org/sigclip_vision_384/resolve/main/sigclip_vision_patch14_384.safetensors" \
   "$MODELS_DIR/clip_vision/sigclip_vision_patch14_384.safetensors"

# ControlNet
log "[ ControlNet ]"
dl "https://huggingface.co/Shakker-Labs/FLUX.1-dev-ControlNet-Union-Pro/resolve/main/diffusion_pytorch_model.safetensors" \
   "$MODELS_DIR/controlnet/flux_controlnet_pro.safetensors"

# VAE
log "[ VAE ]"
dl "https://huggingface.co/black-forest-labs/FLUX.1-schnell/resolve/main/ae.safetensors" \
   "$MODELS_DIR/vae/ae.safetensors"

# Upscale
log "[ Upscale ]"
dl "https://huggingface.co/uwg/upscaler/resolve/main/ESRGAN/4x_NMKD-Siax_200k.pth" \
   "$MODELS_DIR/upscale_models/4x_NMKD-Siax_200k.pth"

# PuLID
log "[ PuLID ]"
dl "https://huggingface.co/guozinan/PuLID/resolve/main/pulid_flux_v0.9.0.safetensors" \
   "$MODELS_DIR/pulid/pulid_flux_v0.9.0.safetensors"

# SAM / SAM2
log "[ SAM ]"
dl "https://huggingface.co/Kijai/sam2-safetensors/resolve/main/sam2_hiera_base_plus.safetensors" \
   "$MODELS_DIR/sam2/sam2_hiera_base_plus.safetensors"
dl "https://huggingface.co/datasets/Gourieff/ReActor/resolve/main/models/sams/sam_vit_b_01ec64.pth" \
   "$MODELS_DIR/sams/sam_vit_b_01ec64.pth"

# Facexlib
log "[ Facexlib ]"
dl "https://huggingface.co/leonelhs/facexlib/resolve/main/detection_Resnet50_Final.pth" \
   "$MODELS_DIR/facexlib/detection_Resnet50_Final.pth"
dl "https://huggingface.co/leonelhs/facexlib/resolve/main/parsing_bisenet.pth" \
   "$MODELS_DIR/facexlib/parsing_bisenet.pth"

# InsightFace (antelopev2)
log "[ InsightFace (antelopev2) ]"
mkdir -p "$MODELS_DIR/insightface/models/antelopev2"
dl "https://huggingface.co/DIAMONIK7777/antelopev2/resolve/main/1k3d68.onnx" \
   "$MODELS_DIR/insightface/models/antelopev2/1k3d68.onnx"
dl "https://huggingface.co/DIAMONIK7777/antelopev2/resolve/main/2d106det.onnx" \
   "$MODELS_DIR/insightface/models/antelopev2/2d106det.onnx"
dl "https://huggingface.co/DIAMONIK7777/antelopev2/resolve/main/genderage.onnx" \
   "$MODELS_DIR/insightface/models/antelopev2/genderage.onnx"
dl "https://huggingface.co/DIAMONIK7777/antelopev2/resolve/main/glintr100.onnx" \
   "$MODELS_DIR/insightface/models/antelopev2/glintr100.onnx"
dl "https://huggingface.co/DIAMONIK7777/antelopev2/resolve/main/scrfd_10g_bnkps.onnx" \
   "$MODELS_DIR/insightface/models/antelopev2/scrfd_10g_bnkps.onnx"

# Ultralytics (YOLO)
log "[ Ultralytics ]"
dl "https://huggingface.co/Bingsu/adetailer/resolve/main/face_yolov8m.pt" \
   "$MODELS_DIR/ultralytics/face_yolov8m.pt"
dl "https://huggingface.co/Bingsu/adetailer/resolve/main/hand_yolov8s.pt" \
   "$MODELS_DIR/ultralytics/hand_yolov8s.pt"
dl "https://huggingface.co/Bingsu/adetailer/resolve/main/person_yolov8m-seg.pt" \
   "$MODELS_DIR/ultralytics/person_yolov8m-seg.pt"

# USO (LoRA + projector)
log "[ USO ]"
dl "https://huggingface.co/Comfy-Org/USO_1.0_Repackaged/resolve/main/split_files/loras/uso-flux1-dit-lora-v1.safetensors" \
   "$MODELS_DIR/loras/uso-flux1-dit-lora-v1.safetensors"
dl "https://huggingface.co/Comfy-Org/USO_1.0_Repackaged/resolve/main/split_files/model_patches/uso-flux1-projector-v1.safetensors" \
   "$MODELS_DIR/model_patches/uso-flux1-projector-v1.safetensors"

# Florence-2-base (LLM)
log "[ Florence-2-base ]"
mkdir -p "$MODELS_DIR/LLM/Florence-2-base"
FLORENCE_BASE="https://huggingface.co/microsoft/Florence-2-base/resolve/main"
for f in model.safetensors pytorch_model.bin config.json configuration_florence2.py \
         modeling_florence2.py preprocessor_config.json processing_florence2.py \
         tokenizer.json tokenizer_config.json vocab.json; do
    dl "$FLORENCE_BASE/$f" "$MODELS_DIR/LLM/Florence-2-base/$f"
done

# ── Create Symlinks (matching local layout) ──────────────────
log ""
log "Creating model symlinks..."
symlink "$MODELS_DIR/checkpoints/flux1-dev-fp8.safetensors" "$MODELS_DIR/diffusion_models/flux1-dev-fp8.safetensors"
symlink "$MODELS_DIR/checkpoints/flux1-dev-fp8.safetensors" "$MODELS_DIR/diffusion_models/FLUX/flux1-dev-fp8.safetensors"
symlink "$MODELS_DIR/checkpoints/flux1-dev-fp8.safetensors" "$MODELS_DIR/unet/flux1-dev-fp8.safetensors"
symlink "$MODELS_DIR/controlnet/flux_controlnet_pro.safetensors" "$MODELS_DIR/diffusion_models/flux_controlnet_pro.safetensors"
symlink "$MODELS_DIR/controlnet/flux_controlnet_pro.safetensors" "$MODELS_DIR/xlabs/controlnets/flux_controlnet_pro.safetensors"

# ── Summary ──────────────────────────────────────────────────
echo ""
echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"

if [ ${#FAILED_DOWNLOADS[@]} -gt 0 ]; then
    err "Failed downloads (${#FAILED_DOWNLOADS[@]}):"
    for f in "${FAILED_DOWNLOADS[@]}"; do
        err "  - $f"
    done
    echo ""
fi

if [ ${#FAILED_NODES[@]} -gt 0 ]; then
    warn "Custom nodes with dependency issues (${#FAILED_NODES[@]}):"
    for n in "${FAILED_NODES[@]}"; do
        warn "  - $n"
    done
    echo ""
fi

if [ ${#FAILED_DOWNLOADS[@]} -eq 0 ] && [ ${#FAILED_NODES[@]} -eq 0 ]; then
    log "Setup complete — no errors!"
else
    log "Setup complete with issues (see above)."
fi

echo ""
log "To start ComfyUI:"
log "  cd $SCRIPT_DIR"
log "  source venv/bin/activate"
log "  python main.py --listen 0.0.0.0 --lowvram --fp8_e4m3fn-text-enc"
echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"