#!/bin/bash
# Setup script for Stable Diffusion WebUI Forge
# Target: Ubuntu 22.04 LTS + NVIDIA GPU (tested on RTX 2000 Ada 16GB)

set -e

SD_DIR="$HOME/stable-diffusion-webui"

echo "=== Stable Diffusion WebUI Forge Setup ==="
echo "=== Target: Ubuntu 22.04 LTS + NVIDIA GPU ==="
echo ""

# 1. Install system dependencies that Ubuntu 22.04 doesn't ship by default
echo "Installing system dependencies..."
sudo apt-get update
sudo apt-get install -y \
    git \
    wget \
    python3-venv \
    python3-pip \
    libgl1 \
    libglib2.0-0 \
    libsm6 \
    libxrender1 \
    libxext6 \
    libffi-dev \
    libssl-dev

# 2. Check NVIDIA driver
if ! command -v nvidia-smi &>/dev/null; then
    echo ""
    echo "ERROR: nvidia-smi not found. Install NVIDIA drivers first:"
    echo ""
    echo "  sudo apt install nvidia-driver-535"
    echo "  sudo reboot"
    echo ""
    echo "Or use Ubuntu's driver manager:"
    echo "  sudo ubuntu-drivers autoinstall"
    echo "  sudo reboot"
    exit 1
fi
echo ""
echo "GPU detected:"
nvidia-smi --query-gpu=name,driver_version,memory.total --format=csv,noheader
echo ""

# 3. Check CUDA availability
if ! command -v nvcc &>/dev/null; then
    echo "NOTE: nvcc not found — Forge will download its own CUDA via PyTorch, which is fine."
    echo "      (No need to install CUDA toolkit separately.)"
    echo ""
fi

# 4. Clone Forge (AUTOMATIC1111 fork with better performance)
if [ -d "$SD_DIR" ]; then
    echo "Directory $SD_DIR already exists, pulling latest..."
    cd "$SD_DIR" && git pull
else
    echo "Cloning Stable Diffusion WebUI Forge..."
    git clone https://github.com/lllyasviel/stable-diffusion-webui-forge.git "$SD_DIR"
fi
cd "$SD_DIR"

# 5. Download models
MODEL_DIR="$SD_DIR/models/Stable-diffusion"
mkdir -p "$MODEL_DIR"

# SDXL Base 1.0 — default model (~6.9GB)
if [ ! -f "$MODEL_DIR/sd_xl_base_1.0.safetensors" ]; then
    echo "Downloading SDXL Base 1.0 model (~6.9GB)..."
    wget -q --show-progress -O "$MODEL_DIR/sd_xl_base_1.0.safetensors" \
        "https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0/resolve/main/sd_xl_base_1.0.safetensors"
else
    echo "SDXL Base 1.0 already downloaded."
fi

# Juggernaut XL v9 — uncensored model (~6.6GB)
if [ ! -f "$MODEL_DIR/juggernautXL_v9.safetensors" ]; then
    echo "Downloading Juggernaut XL v9 model (~6.6GB)..."
    wget -q --show-progress -O "$MODEL_DIR/juggernautXL_v9.safetensors" \
        "https://huggingface.co/RunDiffusion/Juggernaut-XL-v9/resolve/main/Juggernaut-XL_v9_RunDiffusionPhoto_v2.safetensors"
else
    echo "Juggernaut XL v9 already downloaded."
fi

# 6. First launch to create venv (will likely fail on CLIP — that's expected)
echo ""
echo "Running first launch to create Python venv..."
echo "(CLIP build failure is expected — we'll fix it next)"
echo ""
./webui.sh --api --listen --xformers --no-half-vae --exit || true

# 7. Fix CLIP build issue (Ubuntu 22.04 / Python 3.10)
if [ -d "$SD_DIR/venv" ]; then
    echo ""
    echo "Fixing CLIP build dependencies..."
    venv/bin/pip install "setuptools<70" wheel
    venv/bin/pip install --no-build-isolation \
        https://github.com/openai/CLIP/archive/d50d76daa670286dd6cacf3bcd80b5e4823fc8e1.zip
fi

echo ""
echo "=== Setup complete ==="
echo ""
echo "To start the WebUI with API enabled, run:"
echo "  cd $SD_DIR"
echo "  ./webui.sh --api --listen --xformers --no-half-vae"
echo ""
echo "First full launch will take several minutes (installs PyTorch, xformers, etc.)"
echo ""
echo "The API will be available at http://localhost:7860"
echo "Test it with: curl http://localhost:7860/sdapi/v1/sd-models"
echo ""
echo "To select the default SDXL model via API:"
echo "  curl -X POST http://localhost:7860/sdapi/v1/options \\"
echo "      -H 'Content-Type: application/json' \\"
echo "      -d '{\"sd_model_checkpoint\": \"sd_xl_base_1.0.safetensors\"}'"
echo ""
echo "To run as a systemd service:"
echo "  chmod +x setup-sd-service.sh && sudo ./setup-sd-service.sh"