#!/usr/bin/env bash

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

    # Settings
    TARGET_DIR="${1:-$HOME/sd-webui}"
    PYTHON_BIN="${PYTHON_BIN:-python3}"
    REPO_URL="https://github.com/AUTOMATIC1111/stable-diffusion-webui.git"
    ENABLE_XFORMERS="${ENABLE_XFORMERS:-auto}"   # auto|yes|no

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

    if command -v apt >/dev/null 2>&1; then
      echo "[STEP] Installing dependencies (APT)"
      sudo apt-get update -y
      sudo apt-get install -y git ${PYTHON_BIN} ${PYTHON_BIN}-venv build-essential ffmpeg
    fi

    # GPU detection
    GPU_VENDOR="cpu"
    if command -v nvidia-smi >/dev/null 2>&1; then
      GPU_VENDOR="nvidia"
      echo "[INFO] Detected NVIDIA GPU"
      nvidia-smi || true
    else
      if lspci | grep -qi 'amd/ati'; then GPU_VENDOR="amd"; fi
      if lspci | grep -qi 'intel'; then GPU_VENDOR="intel"; fi
      echo "[INFO] Detected GPU vendor: $GPU_VENDOR"
    fi

    if [[ ! -d stable-diffusion-webui ]]; then
      echo "[STEP] Cloning AUTOMATIC1111 webui"
      git clone --depth=1 "$REPO_URL"
    fi
    cd stable-diffusion-webui

    echo "[STEP] Creating Python venv"
    ${PYTHON_BIN} -m venv .venv
    source .venv/bin/activate

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

    # Optional performance packages
    if [[ "$GPU_VENDOR" == "nvidia" ]]; then
      if [[ "$ENABLE_XFORMERS" == "yes" || "$ENABLE_XFORMERS" == "auto" ]]; then
        echo "[STEP] Installing xformers (optional)"
        pip install --extra-index-url https://download.pytorch.org/whl/cu121 xformers || true
      fi
    fi

    echo "[TIP] Place your models in: $TARGET_DIR/stable-diffusion-webui/models/Stable-diffusion"
    echo "[RUN] To start: cd '$TARGET_DIR/stable-diffusion-webui' && source .venv/bin/activate && \
python launch.py --xformers --opt-sdp-attention"