#!/bin/bash
set -e

echo "🚀 Setting up file structure and repositories..."

# Create directory structure
mkdir -p comfyui/ComfyUI
mkdir -p workflows models/checkpoints models/loras models/vae outputs temp

# Define ComfyUI paths for clarity
COMFYUI_DIR="comfyui/ComfyUI"
CUSTOM_NODES_DIR="$COMFYUI_DIR/custom_nodes"

# Clone ComfyUI if not exists
if [ ! -d "$COMFYUI_DIR/.git" ]; then
    echo "📦 Cloning ComfyUI..."
    git clone --depth 1 https://github.com/comfyanonymous/ComfyUI.git "$COMFYUI_DIR"
fi

# Install ComfyUI's own Python dependencies, including the frontend
echo "📦 Installing ComfyUI's Python requirements..."
pip install -r "$COMFYUI_DIR/requirements.txt"

# Install custom nodes
echo "📦 Installing custom nodes..."
cd "$CUSTOM_NODES_DIR"

if [ ! -d "ComfyUI-Impact-Pack" ]; then
    git clone --depth 1 https://github.com/ltdrdata/ComfyUI-Impact-Pack.git
fi
if [ ! -d "facerestore_cf" ]; then
    git clone --depth 1 https://github.com/mav-rik/facerestore_cf.git
fi

cd ../../.. # Return to the root directory of the app

# Download models required by custom nodes
echo "📦 Downloading models for custom nodes..."

FACERESTORE_MODEL_DIR="$COMFYUI_DIR/models/facerestore"
GFPGAN_MODEL_PATH="$FACERESTORE_MODEL_DIR/GFPGANv1.4.pth"

# Create the directory if it doesn't exist
mkdir -p "$FACERESTORE_MODEL_DIR"

# Download GFPGAN model if it's missing
if [ ! -f "$GFPGAN_MODEL_PATH" ]; then
    echo "📥 Downloading GFPGANv1.4 model..."
    wget -q -O "$GFPGAN_MODEL_PATH" https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.4.pth
    echo "✅ GFPGANv1.4 model downloaded."
else
    echo "✅ Found existing GFPGANv1.4 model."
fi

echo "✅ Setup script finished."