#!/usr/bin/env bash

set -euo pipefail
IFS=$'\n\t'
trap 'echo "[ERROR] Line ${LINENO}: command failed" >&2' ERR

# Settings – change if you want a different folder
TARGET_DIR="${1:-$HOME/sd-webui}"
PYTHON_BIN="${PYTHON_BIN:-python3}"
REPO_URL="https://github.com/AUTOMATIC1111/stable-diffusion-webui.git"

echo "[INFO] Target install directory: $TARGET_DIR"

echo "[STEP] Updating APT & installing system dependencies"
sudo apt update
sudo apt install -y \
  git \
  "$PYTHON_BIN"-venv \
  "$PYTHON_BIN"-dev \
  python3-pip \
  ffmpeg \
  libgl1 \
  libglib2.0-0

echo "[STEP] Checking NVIDIA GPU availability"
if command -v nvidia-smi >/dev/null 2>&1; then
  nvidia-smi || true
else
  echo "[WARN] nvidia-smi not found. Make sure NVIDIA driver is installed and you rebooted." >&2
fi

echo "[STEP] Creating target directory: $TARGET_DIR"
mkdir -p "$TARGET_DIR"
cd "$TARGET_DIR"

if [[ ! -d stable-diffusion-webui ]]; then
  echo "[STEP] Cloning AUTOMATIC1111 Stable Diffusion WebUI"
  git clone "$REPO_URL"
else
  echo "[STEP] Repo already exists; pulling latest changes"
  cd stable-diffusion-webui
  git pull
  cd ..
fi

cd stable-diffusion-webui

echo "[STEP] Creating Python virtual environment (.venv) if missing"
if [[ ! -d .venv ]]; then
  "$PYTHON_BIN" -m venv .venv
fi

echo "[STEP] Activating virtual environment"
# shellcheck disable=SC1091
source .venv/bin/activate

echo "[STEP] Upgrading pip / wheel / setuptools"
pip install --upgrade pip wheel setuptools

echo "[STEP] Installing CUDA-enabled PyTorch (12.1 wheels)"
# This uses the official CUDA 12.1 wheel index
pip install --index-url https://download.pytorch.org/whl/cu121 \
  torch torchvision torchaudio

echo "[STEP] Installing xformers (for faster attention on NVIDIA GPUs)"
pip install --extra-index-url https://download.pytorch.org/whl/cu121 xformers || {
  echo "[WARN] xformers installation failed. WebUI will still run, just a bit slower." >&2
}

echo "[STEP] Letting WebUI manage its own Python deps on first run"
# (AUTOMATIC1111 will install the rest of the requirements when you launch)

echo
echo "[DONE] Stable Diffusion WebUI setup finished."

cat <<'EOM'

===============================================
How to run Stable Diffusion WebUI (RTX 3060)
===============================================

1. Activate the venv and go to the repo directory:

   cd ~/sd-webui/stable-diffusion-webui
   source .venv/bin/activate

2. Export recommended env vars for RTX 3060:

   export TORCH_CUDA_ARCH_LIST="8.6"
   export COMMANDLINE_ARGS="--xformers --opt-sdp-attention"

3. Start the WebUI:

   python launch.py

4. Open your browser and go to:

   http://<server-ip>:7860/

Place your .ckpt / .safetensors models in:

   ~/sd-webui/stable-diffusion-webui/models/Stable-diffusion

EOM