#!/usr/bin/env bash
#
#############################################################################################################################
#
# Wael Isa
# Website:  https://www.wael.name
# GitHub:   https://github.com/waelisa
# Version:  v1.1.6
# Build Date: 03/19/2026
# License: MIT
#
# ██╗    ██╗ █████╗ ███████╗██╗         ██╗███████╗ █████╗
# ██║    ██║██╔══██╗██╔════╝██║         ██║██╔════╝██╔══██╗
# ██║ █╗ ██║███████║█████╗  ██║         ██║███████╗███████║
# ██║███╗██║██╔══██║██╔══╝  ██║         ██║╚════██║██╔══██║
# ╚███╔███╔╝██║  ██║███████╗███████╗    ██║███████╗██║  ██║
# ╚══╝╚══╝ ╚═╝  ╚═╝╚══════╝╚══════╝    ╚═╝╚══════╝╚═╝  ╚═╝
#
# Description:
#   Interactive menu to install and manage Stable Diffusion WebUI (AUTOMATIC1111), Forge, and ComfyUI.
#   - Uses gum for a beautiful menu.
#   - Automatically detects GPU (NVIDIA, AMD, Intel) and installs correct PyTorch backend.
#   - Sets optimal launch arguments based on VRAM and GPU type (best‑practices tuning).
#   - Shares models between WebUIs via symlinks (A1111/Forge) or extra_model_paths.yaml (ComfyUI).
#   - Installs ComfyUI‑Manager and its dependencies.
#   - Enforces Python 3.11 with fallback to 3.10/3.12 if needed, with future‑proof checks for Python 3.13+.
#   - Distro-agnostic: detects package manager and can install system dependencies.
#   - Includes “Fix Paths” utility to repair broken symlinks and configs after moving folders.
#   - Adds ffmpeg for video workflows.
#   - Applies modern arguments: --theme=dark, --no-half-vae, and --opt-sdp-no-mem-attention.
#   - Security: verifies repository URLs and directory permissions.
#   - Caches hardware info to avoid repeated detection.
#   - Verbose mode with timestamps and debug logs.
#   - GPU temperature monitoring with thermal warnings.
#   - Confirmation before resetting virtual environments.
#   - Automated dependency updates in non‑interactive mode.
#   - Unified model hub (optional) to save even more space.
#   - Pre‑flight safety checks (disk space, pacman lock, internet).
#   - Hardware dashboard with temperature, VRAM usage, and system RAM.
#   - Architecture‑aware model linking (SD1.5, SDXL, Flux).
#   - Automatic system cleanup after successful installation.
#   - Continuous GPU health monitor (thermal throttle warnings).
#   - Dependency conflict checker (onnxruntime, torch versions).
#   - One‑click installation of recommended "dvine" model from CivitAI.
#   - Multi‑model installer for BoRnNo0b’s checkpoint collection.
#   - Post‑install prompt to install models.
#
# Features:
#   • Install AUTOMATIC1111 WebUI
#   • Install Forge WebUI (uses requirements_versions.txt)
#   • Install ComfyUI (with Manager)
#   • Install multiple WebUIs with model sharing
#   • Link models between A1111/Forge (symlinks)
#   • Link models to ComfyUI (extra_model_paths.yaml)
#   • Set up unified model hub (central storage)
#   • Install recommended model (dvine)
#   • Install multiple models from BoRnNo0b’s collection
#   • Auto‑optimize launch configs for all installed UIs
#   • Fix Paths (repair broken links/configs)
#   • Toggle attention mechanism (xformers vs sdp-no-mem)
#   • Hardware status dashboard
#   • Continuous GPU health monitor
#   • Check for dependency conflicts
#   • Update all dependencies (pip upgrade)
#   • Reset venv for any repository (with confirmation)
#   • Manual PyTorch installation for specific GPU types
#   • Install system dependencies
#   • Verbose mode (--verbose flag)
#
# Changelog:
#   v1.1.6 - Configure webui-user.sh with python_cmd, TORCH_COMMAND, and REQS_FILE.
#            This ensures correct Python and PyTorch are used even after venv reset.
#   v1.1.5 - Fixed pip installation inside virtual environment (externally-managed-environment error).
#   v1.1.4 - Fixed all remaining silent exits (AMD detection, temperature retrieval, cache writes).
#   v1.1.3 - Fixed package manager failures.
#
#############################################################################################################################

set -euo pipefail

# ---------- Configuration ----------
readonly REPO_A1111="https://github.com/AUTOMATIC1111/stable-diffusion-webui.git"
readonly REPO_FORGE="https://github.com/lllyasviel/stable-diffusion-webui-forge.git"
readonly REPO_COMFY="https://github.com/comfyanonymous/ComfyUI.git"
readonly REPO_COMFY_MANAGER="https://github.com/ltdrdata/ComfyUI-Manager.git"

readonly DIR_A1111="stable-diffusion-webui"
readonly DIR_FORGE="stable-diffusion-webui-forge"
readonly DIR_COMFY="ComfyUI"

readonly VENV_DIR="venv"
readonly REQUIREMENTS_FILE="requirements_versions.txt"  # A1111 and Forge use this

# Central models directory (unified hub)
MODELS_ROOT="${HOME}/sd_models"

# Model information for dvine (keep for backward compatibility)
readonly MODEL_NAME="dvine"
readonly MODEL_VERSION_ID="2795248"  # Latest version ID for dvine v2.0
readonly MODEL_DOWNLOAD_URL="https://civitai.com/api/download/models/${MODEL_VERSION_ID}"
readonly MODEL_FILENAME="dvine_v20.safetensors"

# Associative array of BoRnNo0b models (name -> version ID)
# Fill in the actual version IDs from CivitAI model pages
declare -A BORNNOOB_MODELS=(
    ["Dvine"]="2795248"
    ["NAI"]="CHANGE_ME"           # Replace with actual version ID
    ["Matrix Hentai"]="CHANGE_ME"
    ["Tales Hentai"]="CHANGE_ME"
    ["IL"]="CHANGE_ME"
    ["Pornzilla Hentai"]="CHANGE_ME"
)

# Default subfolder for SD1.5 models (used if architecture not specified)
DEFAULT_MODEL_SUBDIR="Stable-diffusion/SD15"

# Python version preference (fallback order)
readonly PYTHON_VERSIONS=("3.11" "3.10" "3.12")
PYTHON_CMD=""

# Colors for output (fallback when gum not used)
readonly RED='\033[0;31m'
readonly GREEN='\033[0;32m'
readonly YELLOW='\033[1;33m'
readonly BLUE='\033[0;34m'
readonly NC='\033[0m' # No Color

# GPU variables
GPU_TYPE=""        # nvidia, amd, intel, cpu
GPU_VRAM=""        # total VRAM in MB (all vendors)
GPU_GFX=""         # AMD GFX architecture (e.g., gfx1030)
GPU_INFO=""        # raw lspci output
GPU_CACHE_FILE="${HOME}/.cache/sd_manager_gpu.cache"
GPU_TEMP=""        # current GPU temperature (if available)
GPU_COMPUTE_CAP="" # NVIDIA compute capability (e.g., 8.9, 9.0, 12.0)

# Attention mechanism preference (xformers vs sdp-no-mem)
ATTENTION_MODE="sdp-no-mem"   # default can be toggled

# Verbose mode (set by --verbose flag)
VERBOSE=false
DEBUG_LOG="/tmp/sd_manager_$(date +%Y%m%d_%H%M%S).log"

# Thermal monitoring flag (set by menu option)
THERMAL_MONITOR_PID=""

# ---------- Helper Functions ----------
log() {
    local level="$1"
    local msg="$2"
    local timestamp=$(date '+%Y-%m-%d %H:%M:%S')
    echo -e "${timestamp} [${level}] ${msg}" >> "$DEBUG_LOG"
    if [[ "$VERBOSE" == true ]]; then
        echo -e "${timestamp} [${level}] ${msg}"
    fi
}

show_help() {
    cat <<EOF
Usage: $0 [OPTION]

Options:
  -h, --help           Show this help message and exit
  --install-a1111      Non-interactive install of AUTOMATIC1111 WebUI
  --install-forge      Non-interactive install of Forge WebUI
  --install-comfy      Non-interactive install of ComfyUI + Manager
  --install-both       Non-interactive install of A1111 and Forge (linked)
  --install-dvine      Non-interactive install of the dvine model
  --install-bornnoob   Non-interactive install of ALL BoRnNo0b models (requires version IDs set)
  --auto-configure     Auto-detect hardware and configure existing installations
  --verbose            Enable verbose logging (output to console and /tmp/sd_manager_*.log)

Interactive menu (default) when run without options.
Requires 'gum' for the menu interface.
EOF
}

# Parse command-line arguments
if [[ $# -gt 0 ]]; then
    case "$1" in
        -h|--help)
            show_help
            exit 0
            ;;
        --install-a1111)
            NON_INTERACTIVE="a1111"
            ;;
        --install-forge)
            NON_INTERACTIVE="forge"
            ;;
        --install-comfy)
            NON_INTERACTIVE="comfy"
            ;;
        --install-both)
            NON_INTERACTIVE="both"
            ;;
        --install-dvine)
            NON_INTERACTIVE="dvine"
            ;;
        --install-bornnoob)
            NON_INTERACTIVE="bornnoob"
            ;;
        --auto-configure)
            NON_INTERACTIVE="configure"
            ;;
        --verbose)
            VERBOSE=true
            shift
            if [[ $# -gt 0 ]]; then
                case "$1" in
                    --install-a1111) NON_INTERACTIVE="a1111" ;;
                    --install-forge) NON_INTERACTIVE="forge" ;;
                    --install-comfy) NON_INTERACTIVE="comfy" ;;
                    --install-both) NON_INTERACTIVE="both" ;;
                    --install-dvine) NON_INTERACTIVE="dvine" ;;
                    --install-bornnoob) NON_INTERACTIVE="bornnoob" ;;
                    --auto-configure) NON_INTERACTIVE="configure" ;;
                    *) echo "Unknown option: $1"; show_help; exit 1 ;;
                esac
            fi
            ;;
        *)
            echo "Unknown option: $1"
            show_help
            exit 1
            ;;
    esac
fi

log "INFO" "Script started with options: $*"

# Check for gum, exit with installation instructions if missing (only in interactive mode)
check_gum() {
    if [[ -z "${NON_INTERACTIVE:-}" ]] && ! command -v gum &> /dev/null; then
        echo -e "${RED}gum is required for interactive mode but not installed.${NC}"
        echo "Please install it using your package manager or run with --help for non-interactive options."
        if command -v pacman &> /dev/null; then
            echo "  sudo pacman -S gum"
        elif command -v apt-get &> /dev/null; then
            echo "  # Download the latest .deb from https://github.com/charmbracelet/gum/releases"
        elif command -v dnf &> /dev/null; then
            echo "  # Download the latest .rpm from https://github.com/charmbracelet/gum/releases"
        else
            echo "  See https://github.com/charmbracelet/gum#installation"
        fi
        log "ERROR" "gum not installed"
        exit 1
    fi
    log "INFO" "gum check passed"
}

# Pre‑flight safety checks
pre_install_checks() {
    echo -e "${BLUE}Running pre‑flight safety checks...${NC}"
    log "INFO" "Running pre‑flight checks"

    if ! ping -c 1 github.com &> /dev/null; then
        echo -e "${RED}No internet connection to GitHub. Please check your network.${NC}"
        log "ERROR" "Internet connection failed"
        exit 1
    fi

    local free_space=$(df -BG . | awk 'NR==2 {print $4}' | sed 's/G//')
    if [[ "$free_space" -lt 20 ]]; then
        echo -e "${RED}Insufficient disk space: ${free_space}GB free (need at least 20GB).${NC}"
        log "ERROR" "Insufficient disk space: ${free_space}GB"
        exit 1
    fi

    if command -v pacman &> /dev/null; then
        if [[ -f /var/lib/pacman/db.lck ]]; then
            echo -e "${RED}Pacman database is locked (/var/lib/pacman/db.lck exists). Please wait for other package operations to finish.${NC}"
            log "ERROR" "Pacman lock file present"
            exit 1
        fi
    fi

    echo -e "${GREEN}Pre‑flight checks passed.${NC}"
    log "INFO" "Pre‑flight checks passed"
}

# Check directory permissions before cloning
check_dir_permissions() {
    local target_dir="$1"
    local parent_dir=$(dirname "$target_dir")
    if [[ ! -w "$parent_dir" ]]; then
        echo -e "${RED}No write permission in ${parent_dir}. Please check permissions.${NC}"
        log "ERROR" "No write permission in $parent_dir"
        exit 1
    fi
    log "INFO" "Directory permissions OK for $parent_dir"
}

# Detect available Python version with future‑proofing for Python 3.13+
detect_python() {
    # First, try to find exact version in PATH
    for ver in "${PYTHON_VERSIONS[@]}"; do
        if command -v "python${ver}" &> /dev/null; then
            PYTHON_CMD="python${ver}"
            echo -e "${GREEN}Using Python ${ver}${NC}"
            log "INFO" "Using Python $ver"
            return 0
        fi
    done

    # If system Python is too new (≥3.13), try to find AUR python311 on Arch
    if command -v python3 &> /dev/null; then
        local pyver=$(python3 -c 'import sys; print(f"{sys.version_info.major}.{sys.version_info.minor}")')
        if [[ "$pyver" =~ ^3\.1[3-9] ]]; then
            echo -e "${YELLOW}System Python is ${pyver} – WebUIs may not be compatible.${NC}"
            if command -v pacman &> /dev/null; then
                if pacman -Qq python311 2>/dev/null; then
                    echo -e "${GREEN}Found python311 package. Using python3.11.${NC}"
                    PYTHON_CMD="python3.11"
                    return 0
                else
                    echo -e "${YELLOW}Consider installing 'python311' from AUR. Falling back to system python3.${NC}"
                fi
            fi
        fi
        echo -e "${YELLOW}Warning: Using system python3 (version ${pyver}) – SD WebUIs may not be compatible.${NC}"
        PYTHON_CMD="python3"
        log "WARNING" "Using system python3 version $pyver"
        return 0
    fi

    echo -e "${RED}No Python interpreter found. Please install Python 3.10, 3.11, or 3.12.${NC}"
    log "ERROR" "No Python found"
    exit 1
}

# Cache GPU detection to avoid repeated calls
cache_gpu_info() {
    mkdir -p "$(dirname "$GPU_CACHE_FILE")" || true
    cat > "$GPU_CACHE_FILE" <<EOF || true
GPU_TYPE="$GPU_TYPE"
GPU_VRAM="$GPU_VRAM"
GPU_GFX="$GPU_GFX"
GPU_INFO="$GPU_INFO"
GPU_COMPUTE_CAP="$GPU_COMPUTE_CAP"
EOF
    log "INFO" "GPU info cached to $GPU_CACHE_FILE"
}

load_gpu_cache() {
    if [[ -f "$GPU_CACHE_FILE" ]]; then
        source "$GPU_CACHE_FILE"
        echo -e "${BLUE}Loaded GPU info from cache: ${GPU_TYPE} (VRAM: ${GPU_VRAM:-unknown} MB)${NC}"
        log "INFO" "Loaded GPU cache: $GPU_TYPE, VRAM $GPU_VRAM"
        return 0
    fi
    return 1
}

# Get GPU temperature (if available) – always returns 0
get_gpu_temp() {
    local temp=""
    case "$GPU_TYPE" in
        "nvidia")
            if command -v nvidia-smi &> /dev/null; then
                temp=$(nvidia-smi --query-gpu=temperature.gpu --format=csv,noheader,nounits | head -1) || true
                echo "$temp °C"
            else
                echo "N/A (nvidia-smi missing)"
            fi
            ;;
        "amd")
            if command -v rocm-smi &> /dev/null; then
                temp=$(rocm-smi --showtemp 2>/dev/null | grep "GPU[0-9]* temperature" | awk '{print $3}') || true
                if [[ -n "$temp" ]]; then
                    echo "$temp °C"
                else
                    echo "N/A (rocm-smi failed)"
                fi
            else
                echo "N/A (rocm-smi missing)"
            fi
            ;;
        "intel")
            echo "N/A"
            ;;
        *)
            echo "N/A"
            ;;
    esac
    return 0
}

# Get GPU VRAM usage (percentage)
get_gpu_vram_usage() {
    local usage=""
    case "$GPU_TYPE" in
        "nvidia")
            if command -v nvidia-smi &> /dev/null; then
                usage=$(nvidia-smi --query-gpu=memory.used --format=csv,noheader,nounits | head -1) || true
                local total=$GPU_VRAM
                if [[ -n "$usage" && -n "$total" && "$total" -gt 0 ]]; then
                    echo "$((usage * 100 / total))%"
                else
                    echo "N/A"
                fi
            else
                echo "N/A"
            fi
            ;;
        *)
            echo "N/A"
            ;;
    esac
}

# Get system RAM available (in MB)
get_system_ram_available() {
    if command -v free &> /dev/null; then
        free -m | awk 'NR==2 {print $7}' || echo "N/A"
    else
        echo "N/A"
    fi
}

# Show hardware dashboard
show_hardware_dashboard() {
    detect_gpu
    echo -e "${BLUE}=== Hardware Status ===${NC}"
    echo -e "GPU Type: $GPU_TYPE"
    echo -e "GPU Temperature: $(get_gpu_temp)"
    echo -e "GPU VRAM Usage: $(get_gpu_vram_usage)"
    echo -e "System RAM Available: $(get_system_ram_available) MB"

    local temp=$(get_gpu_temp | awk '{print $1}')
    if [[ "$temp" =~ ^[0-9]+$ ]] && [[ "$temp" -gt 85 ]]; then
        echo -e "${RED}WARNING: GPU temperature is critically high (${temp}°C). Check cooling!${NC}"
    fi

    local usage_pct=$(get_gpu_vram_usage | sed 's/%//')
    if [[ "$usage_pct" =~ ^[0-9]+$ ]] && [[ "$usage_pct" -gt 90 ]]; then
        echo -e "${YELLOW}WARNING: VRAM usage is very high (${usage_pct}%). Close some applications.${NC}"
    fi

    local ram_avail=$(get_system_ram_available)
    if [[ "$ram_avail" =~ ^[0-9]+$ ]] && [[ "$ram_avail" -lt 1024 ]]; then
        echo -e "${YELLOW}WARNING: System RAM is low (${ram_avail} MB available). Close some applications.${NC}"
    fi
    log "INFO" "Hardware dashboard displayed"
}

# Continuous GPU health monitor (background or foreground)
start_thermal_monitor() {
    if [[ -n "$THERMAL_MONITOR_PID" ]] && kill -0 "$THERMAL_MONITOR_PID" 2>/dev/null; then
        echo -e "${YELLOW}Thermal monitor already running (PID $THERMAL_MONITOR_PID).${NC}"
        return
    fi

    echo -e "${BLUE}Starting background thermal monitor (checks every 30s). Use 'Stop thermal monitor' to end.${NC}"
    (
        while true; do
            sleep 30
            local temp=$(get_gpu_temp | awk '{print $1}')
            if [[ "$temp" =~ ^[0-9]+$ ]] && [[ "$temp" -gt 85 ]]; then
                notify-send "GPU Thermal Warning" "Temperature reached ${temp}°C" 2>/dev/null || \
                echo -e "${RED}GPU temperature is high: ${temp}°C${NC}" >&2
            fi
        done
    ) &
    THERMAL_MONITOR_PID=$!
    log "INFO" "Thermal monitor started with PID $THERMAL_MONITOR_PID"
}

stop_thermal_monitor() {
    if [[ -n "$THERMAL_MONITOR_PID" ]] && kill -0 "$THERMAL_MONITOR_PID" 2>/dev/null; then
        kill "$THERMAL_MONITOR_PID"
        wait "$THERMAL_MONITOR_PID" 2>/dev/null || true
        THERMAL_MONITOR_PID=""
        echo -e "${GREEN}Thermal monitor stopped.${NC}"
        log "INFO" "Thermal monitor stopped"
    else
        echo -e "${YELLOW}No thermal monitor running.${NC}"
    fi
}

# Dependency conflict checker
check_dependency_conflicts() {
    echo -e "${BLUE}Checking for dependency conflicts in installed WebUIs...${NC}"
    log "INFO" "Running dependency conflict check"
    local found_conflict=false
    for dir in "$DIR_A1111" "$DIR_FORGE" "$DIR_COMFY"; do
        if [[ -d "$dir" && -d "$dir/$VENV_DIR" ]]; then
            echo -e "\n${YELLOW}Checking $dir...${NC}"
            source "$dir/$VENV_DIR/bin/activate"
            if pip list 2>/dev/null | grep -q "onnxruntime" && pip list 2>/dev/null | grep -q "onnxruntime-gpu"; then
                echo -e "${RED}Conflict: Both onnxruntime and onnxruntime-gpu are installed in $dir. This can cause issues.${NC}"
                found_conflict=true
            fi
            local torch_ver=$(pip show torch 2>/dev/null | grep Version | awk '{print $2}')
            if [[ -n "$torch_ver" ]]; then
                echo "  PyTorch version: $torch_ver"
            fi
            deactivate
        fi
    done
    if [[ "$found_conflict" == false ]]; then
        echo -e "${GREEN}No major dependency conflicts found.${NC}"
    fi
    log "INFO" "Dependency conflict check completed"
}

# Detect GPU vendor and capabilities
detect_gpu() {
    if load_gpu_cache; then
        return
    fi

    echo -e "${BLUE}Detecting hardware...${NC}"
    log "INFO" "Starting hardware detection"

    # Get GPU info from lspci, but don't fail if lspci is missing
    if command -v lspci &> /dev/null; then
        GPU_INFO=$(lspci 2>/dev/null | grep -E "VGA compatible controller|Display controller" | head -1) || true
    else
        GPU_INFO=""
        log "WARNING" "lspci not found"
    fi

    log "INFO" "Primary GPU info: $GPU_INFO"

    if [[ -z "$GPU_INFO" ]]; then
        # Fallback: assume CPU mode if no GPU info
        GPU_TYPE="cpu"
        echo -e "${YELLOW}No GPU detected via lspci, falling back to CPU mode.${NC}"
        cache_gpu_info
        return
    fi

    if echo "$GPU_INFO" | grep -iq "nvidia"; then
        GPU_TYPE="nvidia"
        echo -e "${GREEN}NVIDIA GPU detected.${NC}"
        log "INFO" "NVIDIA GPU detected"

        if command -v nvidia-smi &> /dev/null; then
            GPU_VRAM=$(nvidia-smi --query-gpu=memory.total --format=csv,noheader,nounits | head -1) || true
            if [[ -n "$GPU_VRAM" ]]; then
                echo -e "  VRAM: ${GPU_VRAM} MB"
                log "INFO" "VRAM: $GPU_VRAM MB"
            fi

            GPU_COMPUTE_CAP=$(nvidia-smi --query-gpu=compute_cap --format=csv,noheader | head -1 | tr -d '.') || true
            if [[ -n "$GPU_COMPUTE_CAP" ]]; then
                echo -e "  Compute Capability: ${GPU_COMPUTE_CAP}"
                log "INFO" "Compute capability: $GPU_COMPUTE_CAP"

                if [[ "$GPU_COMPUTE_CAP" -ge 120 ]]; then
                    echo -e "${YELLOW}  Note: RTX 50-series (Blackwell) detected. For optimal performance, use PyTorch 2.7+ with CUDA 12.8.${NC}"
                    log "INFO" "RTX 50-series detected, may need PyTorch 2.7+"
                fi
            fi
        else
            echo -e "${YELLOW}  nvidia-smi not found. Install NVIDIA drivers for VRAM detection.${NC}"
            log "WARNING" "nvidia-smi not found"
        fi
    elif echo "$GPU_INFO" | grep -iq "amd"; then
        GPU_TYPE="amd"
        echo -e "${GREEN}AMD GPU detected.${NC}"
        log "INFO" "AMD GPU detected"

        if command -v rocminfo &> /dev/null; then
            GPU_GFX=$(rocminfo 2>/dev/null | awk '/^Agent 2/,/^$/ {if ($1 == "Name:" && $2 ~ /^gfx/) {print $2; exit}}') || true
            if [[ -n "$GPU_GFX" ]]; then
                echo -e "  GFX architecture: ${GPU_GFX}"
                log "INFO" "GFX architecture: $GPU_GFX"
                if [[ "$GPU_GFX" == gfx9* ]]; then
                    export HSA_OVERRIDE_GFX_VERSION=9.0.0
                    echo -e "${YELLOW}  Setting HSA_OVERRIDE_GFX_VERSION=9.0.0 for compatibility${NC}"
                fi
            fi
        else
            echo -e "${YELLOW}  rocminfo not found. Install ROCm for detailed detection.${NC}"
            log "WARNING" "rocminfo not found"
        fi

        if [[ -d /sys/class/drm/ ]]; then
            for card in /sys/class/drm/card?/device/mem_info_vram_total; do
                if [[ -f "$card" ]]; then
                    vram_bytes=$(cat "$card" 2>/dev/null) || true
                    if [[ -n "$vram_bytes" ]]; then
                        GPU_VRAM=$((vram_bytes / 1024 / 1024))
                        echo -e "  VRAM: ${GPU_VRAM} MB"
                        log "INFO" "VRAM: $GPU_VRAM MB (from sysfs)"
                    fi
                    break
                fi
            done
        fi
        if [[ -z "$GPU_VRAM" ]] && command -v glxinfo &> /dev/null; then
            GPU_VRAM=$(glxinfo -B 2>/dev/null | awk '/Video memory:/ {print $3}') || true
            if [[ -n "$GPU_VRAM" ]]; then
                echo -e "  VRAM: ${GPU_VRAM} MB"
                log "INFO" "VRAM: $GPU_VRAM MB (from glxinfo)"
            fi
        fi
    elif echo "$GPU_INFO" | grep -iq "intel"; then
        GPU_TYPE="intel"
        echo -e "${GREEN}Intel GPU detected.${NC}"
        log "INFO" "Intel GPU detected"

        if [[ -d /sys/class/drm/ ]]; then
            for card in /sys/class/drm/card?/device/mem_info_vram_total; do
                if [[ -f "$card" ]]; then
                    vram_bytes=$(cat "$card" 2>/dev/null) || true
                    if [[ -n "$vram_bytes" ]]; then
                        GPU_VRAM=$((vram_bytes / 1024 / 1024))
                        echo -e "  VRAM: ${GPU_VRAM} MB"
                        log "INFO" "VRAM: $GPU_VRAM MB (from sysfs)"
                    fi
                    break
                fi
            done
        fi
        if [[ -z "$GPU_VRAM" ]] && command -v glxinfo &> /dev/null; then
            GPU_VRAM=$(glxinfo -B 2>/dev/null | awk '/Video memory:/ {print $3}') || true
            if [[ -n "$GPU_VRAM" ]]; then
                echo -e "  VRAM: ${GPU_VRAM} MB"
                log "INFO" "VRAM: $GPU_VRAM MB (from glxinfo)"
            fi
        fi
    else
        GPU_TYPE="cpu"
        echo -e "${YELLOW}No supported GPU found, will use CPU mode.${NC}"
        log "WARNING" "No supported GPU, falling back to CPU"
    fi

    GPU_TEMP=$(get_gpu_temp)
    echo -e "  Temperature: $GPU_TEMP"
    log "INFO" "GPU temperature: $GPU_TEMP"

    cache_gpu_info
}

# Generate the correct PyTorch pip command based on GPU type
get_torch_command() {
    local target_gpu="${1:-$GPU_TYPE}"
    case "$target_gpu" in
        "nvidia")
            if [[ -n "$GPU_COMPUTE_CAP" && "$GPU_COMPUTE_CAP" -ge 120 ]]; then
                echo "pip install torch torchvision torchaudio xformers --index-url https://download.pytorch.org/whl/cu128"
            else
                echo "pip install torch torchvision torchaudio xformers --index-url https://download.pytorch.org/whl/cu124"
            fi
            ;;
        "amd")
            if [[ -n "$GPU_GFX" ]]; then
                if [[ "$GPU_GFX" == gfx9* ]]; then
                    echo "pip install torch==1.13.1+rocm5.2 torchvision==0.14.1+rocm5.2 --index-url https://download.pytorch.org/whl/rocm5.2"
                elif [[ "$GPU_GFX" == gfx10* ]]; then
                    echo "pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/rocm5.7"
                elif [[ "$GPU_GFX" == gfx11* ]]; then
                    echo "pip install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/rocm6.0"
                else
                    echo "pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/rocm7.1"
                fi
            else
                echo "pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/rocm7.1"
            fi
            ;;
        "intel")
            echo "pip install torch==2.1.0 torchvision==0.16.0 torchaudio==2.1.0 --index-url https://download.pytorch.org/whl/xpu && pip install intel-extension-for-pytorch==2.1.10+xpu --extra-index-url https://pytorch-extension.intel.com/release-whl/stable/xpu/us/"
            ;;
        *)
            echo "pip install torch torchvision torchaudio"
            ;;
    esac
}

# Configure webui-user.sh for A1111/Forge
configure_webui_user_sh() {
    local target_dir="$1"
    local config_file="${target_dir}/webui-user.sh"

    if [[ ! -f "$config_file" ]]; then
        # Create a default if missing (should exist after clone, but just in case)
        cat > "$config_file" <<'EOF'
#!/bin/bash
################################################################################
### AUTOMATIC1111 Stable Diffusion WebUI Launch Script
################################################################################

# Command line arguments
export COMMANDLINE_ARGS=""

# Python executable (optional)
# export PYTHON=""

# Git executable (optional)
# export GIT=""

# Virtual environment directory (optional)
# export VENV_DIR="venv"

# Requirements file (optional)
# export REQS_FILE="requirements_versions.txt"

# PyTorch installation command (optional)
# export TORCH_COMMAND=""

################################################################################
exec "${0/sh/bat}" "$@"
EOF
    fi

    local backup="${config_file}.bak.$(date +%Y%m%d_%H%M%S)"
    cp "$config_file" "$backup"
    log "INFO" "Backup of webui-user.sh created at $backup"

    # Set python_cmd (the variable name is `python_cmd`, not `PYTHON`)
    if grep -q "^python_cmd=" "$config_file"; then
        sed -i "s|^python_cmd=.*|python_cmd=\"$PYTHON_CMD\"|" "$config_file"
    else
        # Insert after shebang or at top
        sed -i "2i python_cmd=\"$PYTHON_CMD\"" "$config_file"
    fi

    # Set TORCH_COMMAND
    local torch_cmd=$(get_torch_command)
    if grep -q "^export TORCH_COMMAND=" "$config_file"; then
        sed -i "s|^export TORCH_COMMAND=.*|export TORCH_COMMAND=\"$torch_cmd\"|" "$config_file"
    else
        echo "export TORCH_COMMAND=\"$torch_cmd\"" >> "$config_file"
    fi

    # Set REQS_FILE
    if grep -q "^export REQS_FILE=" "$config_file"; then
        sed -i "s|^export REQS_FILE=.*|export REQS_FILE=\"$REQUIREMENTS_FILE\"|" "$config_file"
    else
        echo "export REQS_FILE=\"$REQUIREMENTS_FILE\"" >> "$config_file"
    fi

    log "INFO" "Configured $config_file with python_cmd=$PYTHON_CMD, TORCH_COMMAND, REQS_FILE"
}

# Detect package manager and install system dependencies
install_system_deps() {
    echo -e "${YELLOW}Installing system dependencies (Python, git, wget, pciutils, coreutils, ffmpeg)...${NC}"
    log "INFO" "Installing system dependencies"

    if command -v pacman &> /dev/null; then
        if ! sudo pacman -S --needed python python-pip python-virtualenv git wget pciutils coreutils ffmpeg; then
            echo -e "${RED}Failed to install dependencies with pacman. Please check your internet connection and package database.${NC}"
            log "ERROR" "pacman installation failed"
            exit 1
        fi
    elif command -v apt-get &> /dev/null; then
        sudo apt-get update || { echo -e "${RED}apt-get update failed.${NC}"; log "ERROR" "apt-get update failed"; exit 1; }
        if ! sudo apt-get install -y git wget pciutils coreutils ffmpeg; then
            echo -e "${RED}Failed to install base dependencies with apt-get.${NC}"
            log "ERROR" "apt-get install failed"
            exit 1
        fi
        if ! grep -q ^deb .*deadsnakes /etc/apt/sources.list /etc/apt/sources.list.d/* 2>/dev/null; then
            sudo add-apt-repository -y ppa:deadsnakes/ppa || { echo -e "${RED}Failed to add deadsnakes PPA.${NC}"; log "ERROR" "add-apt-repository failed"; exit 1; }
            sudo apt-get update || { echo -e "${RED}apt-get update after PPA failed.${NC}"; log "ERROR" "apt-get update after PPA failed"; exit 1; }
        fi
        sudo apt-get install -y python3.11 python3.11-venv python3-pip || \
        sudo apt-get install -y python3.10 python3.10-venv python3-pip || \
        sudo apt-get install -y python3 python3-venv python3-pip || { echo -e "${RED}Failed to install Python.${NC}"; log "ERROR" "Python installation failed"; exit 1; }
    elif command -v dnf &> /dev/null; then
        sudo dnf makecache --refresh || { echo -e "${RED}dnf makecache failed.${NC}"; log "ERROR" "dnf makecache failed"; exit 1; }
        if ! sudo dnf install -y python3.11 python3-pip git wget pciutils coreutils ffmpeg-free; then
            # Try fallback without specific Python version
            if ! sudo dnf install -y python3 python3-pip git wget pciutils coreutils ffmpeg-free; then
                echo -e "${RED}Failed to install dependencies with dnf.${NC}"
                log "ERROR" "dnf install failed"
                exit 1
            fi
        fi
    else
        echo -e "${RED}Unsupported package manager. Please install dependencies manually:${NC}"
        echo "  - Python 3.10/3.11/3.12"
        echo "  - git"
        echo "  - wget"
        echo "  - pciutils (for hardware detection)"
        echo "  - coreutils (for realpath)"
        echo "  - ffmpeg (for video workflows)"
        log "ERROR" "Unsupported package manager"
        exit 1
    fi

    echo -e "${GREEN}System dependencies installed.${NC}"
    log "INFO" "System dependencies installed"

    if [[ -n "${NON_INTERACTIVE:-}" ]]; then
        echo -e "${YELLOW}Performing system cleanup...${NC}"
        if command -v apt-get &> /dev/null; then
            sudo apt autoremove -y && sudo apt clean
        elif command -v dnf &> /dev/null; then
            sudo dnf autoremove -y && sudo dnf clean all
        elif command -v pacman &> /dev/null; then
            sudo pacman -Sc --noconfirm
        fi
    fi
}

# Create virtual environment
create_venv() {
    local repo_dir="$1"
    local venv_path="${repo_dir}/${VENV_DIR}"
    if [[ -d "$venv_path" ]]; then
        echo -e "${YELLOW}Virtual environment already exists in ${repo_dir}. Skipping creation.${NC}"
        log "INFO" "Venv already exists in $repo_dir"
    else
        echo -e "${GREEN}Creating virtual environment in ${repo_dir} ...${NC}"
        log "INFO" "Creating venv in $repo_dir with $PYTHON_CMD"
        (cd "$repo_dir" && "$PYTHON_CMD" -m venv "$VENV_DIR")
    fi
}

# Install requirements (now runs pip inside the venv)
install_requirements() {
    local repo_dir="$1"
    local req_file="$2"
    local upgrade_flag="${3:-}"

    echo -e "${GREEN}Installing requirements from ${req_file} ...${NC}"
    log "INFO" "Installing requirements from $req_file in $repo_dir"
    (
        cd "$repo_dir" || exit 1
        source "${VENV_DIR}/bin/activate"
        pip install --upgrade pip
        pip install $upgrade_flag -r "$req_file"
    )
}

# Hardware-specific PyTorch installation (now runs pip inside the venv)
install_pytorch() {
    local repo_dir="$1"
    local target_gpu="${2:-$GPU_TYPE}"

    if [[ ! -d "${repo_dir}/${VENV_DIR}" ]]; then
        echo -e "${RED}Virtual environment not found in ${repo_dir}. Please create it first.${NC}"
        log "ERROR" "Venv not found in $repo_dir"
        return 1
    fi

    echo -e "${GREEN}Installing PyTorch for ${target_gpu} GPU in ${repo_dir} ...${NC}"
    log "INFO" "Installing PyTorch for $target_gpu in $repo_dir"

    (
        cd "$repo_dir" || exit 1
        source "${VENV_DIR}/bin/activate"
        eval "$(get_torch_command "$target_gpu")"
    )

    log "INFO" "PyTorch installation completed in $repo_dir"
}

# Verify Git remote URL
verify_repo_url() {
    local repo_url="$1"
    local repo_dir="$2"
    if [[ -d "$repo_dir/.git" ]]; then
        local remote_url=$(git -C "$repo_dir" config --get remote.origin.url)
        local norm_repo="${repo_url%.git}"
        local norm_remote="${remote_url%.git}"
        if [[ "$norm_remote" != "$norm_repo" ]]; then
            echo -e "${RED}Security warning: existing repository remote URL (${remote_url}) does not match expected (${repo_url}).${NC}"
            echo "Please check the repository manually."
            log "ERROR" "Repository URL mismatch for $repo_dir"
            exit 1
        fi
        log "INFO" "Repository URL verified for $repo_dir"
    fi
}

# Clone or update a repository
clone_repo() {
    local repo_url="$1"
    local repo_dir="$2"
    check_dir_permissions "$repo_dir"
    if [[ -d "$repo_dir" ]]; then
        verify_repo_url "$repo_url" "$repo_dir"
        echo -e "${YELLOW}Repository ${repo_dir} already exists. Pulling latest changes...${NC}"
        log "INFO" "Pulling latest changes for $repo_dir"
        (cd "$repo_dir" && git pull)
    else
        echo -e "${GREEN}Cloning ${repo_url} ...${NC}"
        log "INFO" "Cloning $repo_url into $repo_dir"
        git clone "$repo_url" "$repo_dir"
    fi
}

# Reset venv
reset_venv() {
    local repo_dir="$1"
    local venv_path="${repo_dir}/${VENV_DIR}"
    if [[ -d "$venv_path" ]]; then
        if [[ -z "${NON_INTERACTIVE:-}" ]]; then
            gum confirm "Are you sure you want to reset the virtual environment in ${repo_dir}? This will delete the current venv." || return
        fi
        echo -e "${YELLOW}Removing existing virtual environment in ${repo_dir} ...${NC}"
        log "INFO" "Removing venv in $repo_dir"
        rm -rf "$venv_path"
    fi
    create_venv "$repo_dir"
    echo -e "${GREEN}Virtual environment reset complete. You can now install PyTorch and requirements.${NC}"
    log "INFO" "Venv reset in $repo_dir"
}

# ---------- Unified Model Hub Setup ----------
setup_central_models() {
    echo -e "${BLUE}Setting up unified model hub at ${MODELS_ROOT}...${NC}"
    mkdir -p "$MODELS_ROOT"/{Stable-diffusion/SD15,Stable-diffusion/SDXL,Stable-diffusion/Flux,Lora,VAE,embeddings,ControlNet,ESRGAN,RealESRGAN,SwinIR,GFPGAN}

    echo -e "${YELLOW}Do you want to move existing models from your installations to the hub?${NC}"
    if gum confirm "This will move (not copy) model files to $MODELS_ROOT and symlink them back. Proceed?"; then
        for dir in "$DIR_A1111" "$DIR_FORGE" "$DIR_COMFY"; do
            if [[ -d "$dir" ]]; then
                # Handle main model folders with architecture subfolders
                for sub in "models/Stable-diffusion" "models/Lora" "models/VAE" "embeddings" "models/ControlNet" "models/ESRGAN" "models/RealESRGAN" "models/SwinIR" "models/GFPGAN"; do
                    local src="$dir/$sub"
                    local dst_base=$(basename "$sub")
                    if [[ "$sub" == models/* ]]; then
                        dst_base="${sub#models/}"
                    else
                        dst_base="$sub"
                    fi
                    local dst="$MODELS_ROOT/$dst_base"
                    if [[ -d "$src" && ! -L "$src" ]]; then
                        echo "Moving $src to $dst..."
                        mkdir -p "$dst"
                        # For Stable-diffusion, we need to move architecture subfolders individually
                        if [[ "$dst_base" == "Stable-diffusion" ]]; then
                            for arch in "SD15" "SDXL" "Flux"; do
                                if [[ -d "$src/$arch" ]]; then
                                    mv "$src/$arch"/* "$dst/$arch"/ 2>/dev/null || true
                                    rmdir "$src/$arch" 2>/dev/null || true
                                fi
                            done
                            # Move any loose files into SD15 as default
                            find "$src" -maxdepth 1 -type f -name "*.safetensors" -o -name "*.ckpt" -exec mv {} "$dst/SD15"/ \;
                        else
                            mv "$src"/* "$dst"/ 2>/dev/null || true
                        fi
                        rm -rf "$src"
                        ln -s "$dst" "$src"
                    fi
                done
            fi
        done
        echo -e "${GREEN}Models moved and symlinked.${NC}"
    fi

    for ui in "$DIR_A1111" "$DIR_FORGE"; do
        if [[ -d "$ui" ]]; then
            local config="$ui/webui-user.sh"
            if [[ -f "$config" ]]; then
                cp "$config" "$config.bak.central"
                local current_args=""
                if grep -q "^export COMMANDLINE_ARGS=" "$config"; then
                    current_args=$(grep "^export COMMANDLINE_ARGS=" "$config" | sed 's/^export COMMANDLINE_ARGS=//' | sed 's/"//g')
                fi
                local new_args="$current_args"
                for arch in "SD15" "SDXL" "Flux"; do
                    new_args="$new_args --ckpt-dir $MODELS_ROOT/Stable-diffusion/$arch"
                done
                new_args="$new_args --lora-dir $MODELS_ROOT/Lora --vae-dir $MODELS_ROOT/VAE --embeddings-dir $MODELS_ROOT/embeddings"
                if grep -q "^export COMMANDLINE_ARGS=" "$config"; then
                    sed -i "s|^export COMMANDLINE_ARGS=.*|export COMMANDLINE_ARGS=\"${new_args}\"|" "$config"
                else
                    echo "export COMMANDLINE_ARGS=\"${new_args}\"" >> "$config"
                fi
            fi
        fi
    done

    if [[ -d "$DIR_COMFY" ]]; then
        local comfy_config="$DIR_COMFY/extra_model_paths.yaml"
        cat > "$comfy_config" <<EOF
# Unified model hub configuration
hub:
    base_path: $MODELS_ROOT
    checkpoints:
        SD15: Stable-diffusion/SD15
        SDXL: Stable-diffusion/SDXL
        Flux: Stable-diffusion/Flux
    loras: Lora
    vae: VAE
    embeddings: embeddings
    controlnet: ControlNet
    upscale_models: ESRGAN
EOF
    fi

    echo -e "${GREEN}Unified model hub configured.${NC}"
    log "INFO" "Unified hub setup completed"
}

# ---------- Model Linking for A1111/Forge (Architecture-aware) ----------
link_models() {
    local master_dir="$1"
    local target_dir="$2"

    local folders=("models/Stable-diffusion/SD15" "models/Stable-diffusion/SDXL" "models/Stable-diffusion/Flux" "models/Lora" "models/VAE" "embeddings" "models/ControlNet" "models/ESRGAN" "models/RealESRGAN" "models/SwinIR" "models/GFPGAN")

    echo -e "${BLUE}Linking models from ${master_dir} to ${target_dir}...${NC}"
    log "INFO" "Linking models from $master_dir to $target_dir"

    for folder in "${folders[@]}"; do
        local src="${master_dir}/${folder}"
        local dst="${target_dir}/${folder}"

        mkdir -p "$src"

        if [[ -d "$dst" && ! -L "$dst" ]]; then
            echo "Moving existing files from ${dst} to ${src}..."
            log "INFO" "Moving $dst contents to $src"
            cp -rn "$dst"/* "$src"/ 2>/dev/null || true
            rm -rf "$dst"
        fi

        if [[ ! -L "$dst" ]]; then
            ln -s "$(realpath "$src")" "$dst"
            echo -e "${GREEN}Linked: ${folder}${NC}"
            log "INFO" "Linked $folder"
        else
            echo -e "${YELLOW}Link already exists for: ${folder}${NC}"
        fi
    done
}

# ---------- ComfyUI extra_model_paths.yaml generation ----------
generate_extra_model_paths() {
    local comfy_dir="$1"
    local source_ui="$2"
    local source_dir=""

    if [[ "$source_ui" == "A1111" ]]; then
        source_dir="$DIR_A1111"
    elif [[ "$source_ui" == "Forge" ]]; then
        source_dir="$DIR_FORGE"
    else
        echo -e "${RED}Invalid source UI.${NC}"
        log "ERROR" "Invalid source UI: $source_ui"
        return 1
    fi

    if [[ ! -d "$source_dir" ]]; then
        echo -e "${RED}Source directory ${source_dir} does not exist. Please install it first.${NC}"
        log "ERROR" "Source dir $source_dir not found"
        return 1
    fi

    local config_file="${comfy_dir}/extra_model_paths.yaml"
    local base_path
    base_path=$(realpath "$source_dir")

    cat > "$config_file" <<EOF
# This file was auto-generated by sd_manager.sh
# It allows ComfyUI to use models from your existing ${source_ui} installation.

${source_ui,,}:
    base_path: ${base_path}
    checkpoints: models/Stable-diffusion
    configs: models/Stable-diffusion
    vae: models/VAE
    loras: models/Lora
    upscale_models: models/ESRGAN
    embeddings: embeddings
    controlnet: models/ControlNet
EOF
    echo -e "${GREEN}Generated ${config_file} pointing to ${source_ui} models.${NC}"
    log "INFO" "Generated $config_file"
}

# ---------- Fix Paths ----------
fix_paths() {
    echo -e "${BLUE}Checking and fixing paths...${NC}"
    log "INFO" "Running fix_paths"

    if [[ -d "$DIR_A1111" && -d "$DIR_FORGE" ]]; then
        local master=""
        for folder in "models/Stable-diffusion" "models/Lora" "models/VAE"; do
            if [[ -d "$DIR_A1111/$folder" && ! -L "$DIR_A1111/$folder" ]]; then
                master="$DIR_A1111"
                break
            elif [[ -d "$DIR_FORGE/$folder" && ! -L "$DIR_FORGE/$folder" ]]; then
                master="$DIR_FORGE"
                break
            fi
        done
        if [[ -n "$master" ]]; then
            if [[ "$master" == "$DIR_A1111" ]]; then
                echo "Master appears to be AUTOMATIC1111, relinking Forge..."
                link_models "$DIR_A1111" "$DIR_FORGE"
            else
                echo "Master appears to be Forge, relinking AUTOMATIC1111..."
                link_models "$DIR_FORGE" "$DIR_A1111"
            fi
        else
            echo -e "${YELLOW}No clear master found, skipping symlink fix.${NC}"
        fi
    fi

    if [[ -d "$DIR_COMFY" ]]; then
        local config="${DIR_COMFY}/extra_model_paths.yaml"
        if [[ -f "$config" ]]; then
            if grep -q "base_path:" "$config"; then
                local old_path
                old_path=$(grep "base_path:" "$config" | awk '{print $2}')
                if [[ -n "$old_path" && ! -d "$old_path" ]]; then
                    echo -e "${YELLOW}extra_model_paths.yaml points to missing directory: $old_path${NC}"
                    log "WARNING" "extra_model_paths.yaml points to missing $old_path"
                    if [[ -d "$DIR_A1111" ]]; then
                        generate_extra_model_paths "$DIR_COMFY" "A1111"
                    elif [[ -d "$DIR_FORGE" ]]; then
                        generate_extra_model_paths "$DIR_COMFY" "Forge"
                    fi
                fi
            fi
        fi
    fi

    echo -e "${GREEN}Path fixing complete.${NC}"
    log "INFO" "fix_paths completed"
}

# ---------- A1111/Forge Hardware Optimization (COMMANDLINE_ARGS only) ----------
apply_best_tips() {
    local target_dir="$1"
    local config_file="${target_dir}/webui-user.sh"
    local args="--theme=dark --no-half-vae --api --listen"

    echo -e "${BLUE}Optimizing ${config_file} for ${GPU_TYPE} (VRAM: ${GPU_VRAM:-unknown} MB)...${NC}"
    log "INFO" "Optimizing $config_file for $GPU_TYPE"

    case "$GPU_TYPE" in
        "nvidia")
            if [[ "$ATTENTION_MODE" == "xformers" ]]; then
                args="$args --xformers"
            else
                args="$args --opt-sdp-no-mem-attention"
            fi
            args="$args --pin-shared-memory --opt-channelslast"

            if [[ -n "$GPU_VRAM" ]]; then
                if [[ "$GPU_VRAM" -lt 4000 ]]; then
                    args="$args --lowvram"
                elif [[ "$GPU_VRAM" -lt 8000 ]]; then
                    args="$args --medvram"
                elif [[ "$GPU_VRAM" -ge 16000 ]]; then
                    args="$args --highvram"
                    export TORCH_DISTRIBUTED_BACKEND="nccl"
                fi
            fi
            if [[ -n "$GPU_VRAM" && "$GPU_VRAM" -ge 24000 ]]; then
                args="$args --use-uncached-memory"
            fi

            if [[ -n "$GPU_COMPUTE_CAP" && "$GPU_COMPUTE_CAP" -ge 120 ]]; then
                args="$args --opt-sdp-attention"
            fi
            ;;
        "amd")
            args="$args --opt-sdp-no-mem-attention --precision autocast --no-half --opt-sub-quad-attention --disable-nan-check"
            args="$args --use-gpu-memory-pool"
            export TORCH_DISTRIBUTED_BACKEND="gloo"

            if [[ -n "$GPU_GFX" && "$GPU_GFX" == gfx9* ]]; then
                export HSA_OVERRIDE_GFX_VERSION=9.0.0
            fi
            if [[ -n "$GPU_VRAM" && "$GPU_VRAM" -lt 8000 ]]; then
                args="$args --lowvram"
            fi
            ;;
        "intel")
            args="$args --opt-sdp-no-mem-attention --use-ipex --precision full --no-half --opt-sub-quad-attention"
            if [[ "$GPU_INFO" == *"Arc"* ]]; then
                args="$args --lowvram"
            fi
            export TORCH_DISTRIBUTED_BACKEND="gloo"
            ;;
        *)
            args="$args --skip-torch-cuda-test --precision full --no-half --use-cpu all"
            ;;
    esac

    if [[ "$target_dir" == "$DIR_FORGE" ]]; then
        args=$(echo "$args" | sed -E 's/--(low|med|high)vram//g')
        echo -e "${YELLOW}Forge detected: stripped explicit memory flags (Forge uses dynamic memory management).${NC}"
    fi

    args="$args --realtime-monitor"

    if [[ ! -f "$config_file" ]]; then
        # If the file is missing, create it with defaults; later we'll add variables
        cat > "$config_file" <<'EOF'
#!/bin/bash
################################################################################
### AUTOMATIC1111 Stable Diffusion WebUI Launch Script
################################################################################

# Command line arguments
export COMMANDLINE_ARGS=""

# Python executable (optional)
# export PYTHON=""

# Git executable (optional)
# export GIT=""

# Virtual environment directory (optional)
# export VENV_DIR="venv"

# Requirements file (optional)
# export REQS_FILE="requirements_versions.txt"

# PyTorch installation command (optional)
# export TORCH_COMMAND=""

################################################################################
exec "${0/sh/bat}" "$@"
EOF
    fi

    # Backup before modifying
    local backup="${config_file}.bak.$(date +%Y%m%d_%H%M%S)"
    cp "$config_file" "$backup"
    log "INFO" "Backup created at $backup"

    # Update COMMANDLINE_ARGS
    if grep -q "^export COMMANDLINE_ARGS=" "$config_file"; then
        sed -i "s|^export COMMANDLINE_ARGS=.*|export COMMANDLINE_ARGS=\"${args}\"|" "$config_file"
    else
        echo "export COMMANDLINE_ARGS=\"${args}\"" >> "$config_file"
    fi

    echo -e "${GREEN}Applied optimal arguments: ${args}${NC}"
    log "INFO" "Applied args: $args"
}

# ---------- ComfyUI Hardware Optimization ----------
apply_comfy_tips() {
    local target_dir="$1"
    local launch_script="${target_dir}/launch_comfy.sh"
    local args=""

    echo -e "${BLUE}Generating optimized launch script for ComfyUI (${GPU_TYPE})...${NC}"
    log "INFO" "Generating ComfyUI launch script for $GPU_TYPE"

    case "$GPU_TYPE" in
        "nvidia")
            args="--preview-method taesd"
            if [[ -n "$GPU_VRAM" ]]; then
                if [[ "$GPU_VRAM" -lt 4000 ]]; then
                    args="--lowvram $args"
                elif [[ "$GPU_VRAM" -lt 8000 ]]; then
                    args="--normalvram $args"
                elif [[ "$GPU_VRAM" -ge 16000 ]]; then
                    args="--highvram $args"
                fi
            fi
            ;;
        "amd"|"intel")
            args="--use-split-cross-attention --preview-method taesd"
            if [[ -n "$GPU_VRAM" && "$GPU_VRAM" -lt 8000 ]]; then
                args="--lowvram $args"
            fi
            ;;
        *)
            args="--cpu --preview-method taesd"
            ;;
    esac

    cat > "$launch_script" <<EOF
#!/bin/bash
# Optimized launch script for ComfyUI
cd "$(realpath "$target_dir")"
source "${VENV_DIR}/bin/activate"
python main.py $args "\$@"
EOF
    chmod +x "$launch_script"
    echo -e "${GREEN}Created ${launch_script} with arguments: ${args}${NC}"
    echo "You can run ComfyUI with: $launch_script"
    log "INFO" "ComfyUI launch script created with args: $args"
}

# ---------- Toggle attention mechanism ----------
toggle_attention() {
    local current="$ATTENTION_MODE"
    if [[ "$current" == "xformers" ]]; then
        ATTENTION_MODE="sdp-no-mem"
        echo -e "${GREEN}Attention mechanism set to: --opt-sdp-no-mem-attention${NC}"
        log "INFO" "Attention mode set to sdp-no-mem"
    else
        ATTENTION_MODE="xformers"
        echo -e "${GREEN}Attention mechanism set to: --xformers${NC}"
        log "INFO" "Attention mode set to xformers"
    fi
    for dir in "$DIR_A1111" "$DIR_FORGE"; do
        if [[ -d "$dir" ]]; then
            apply_best_tips "$dir"
        fi
    done
}

# ---------- Update all dependencies ----------
update_all_deps() {
    echo -e "${BLUE}Updating dependencies for all installed WebUIs...${NC}"
    log "INFO" "Starting dependency update"
    for dir in "$DIR_A1111" "$DIR_FORGE" "$DIR_COMFY"; do
        if [[ -d "$dir" && -d "$dir/$VENV_DIR" ]]; then
            echo -e "${YELLOW}Updating dependencies in $dir...${NC}"
            log "INFO" "Updating deps in $dir"
            (cd "$dir" && source "$VENV_DIR/bin/activate" && pip install --upgrade pip && pip list --outdated --format=freeze | grep -v '^\-e' | cut -d = -f 1 | xargs -n1 pip install --upgrade || true)
        fi
    done
    echo -e "${GREEN}Dependency update completed.${NC}"
    log "INFO" "Dependency update completed"
}

# ---------- Download a CivitAI model ----------
download_civitai_model() {
    local model_name="$1"
    local version_id="$2"
    local download_url="https://civitai.com/api/download/models/${version_id}"
    local target_dir=""

    # Determine target directory
    if [[ -d "$MODELS_ROOT" ]]; then
        # Use unified hub with architecture subfolder
        target_dir="$MODELS_ROOT/$DEFAULT_MODEL_SUBDIR"
    else
        # Fallback to first installed UI
        local ui_dir=""
        [[ -d "$DIR_A1111" ]] && ui_dir="$DIR_A1111"
        [[ -d "$DIR_FORGE" && -z "$ui_dir" ]] && ui_dir="$DIR_FORGE"
        [[ -d "$DIR_COMFY" && -z "$ui_dir" ]] && ui_dir="$DIR_COMFY"

        if [[ -z "$ui_dir" ]]; then
            echo -e "${RED}No WebUI installation found. Please install a WebUI first.${NC}"
            return 1
        fi

        if [[ "$ui_dir" == "$DIR_COMFY" ]]; then
            target_dir="$ui_dir/models/checkpoints"
        else
            target_dir="$ui_dir/models/Stable-diffusion"
        fi
    fi

    mkdir -p "$target_dir"
    local output_file="${target_dir}/${model_name// /_}.safetensors"

    if [[ -f "$output_file" ]]; then
        echo -e "${YELLOW}Model already exists at $output_file. Skipping.${NC}"
        return
    fi

    echo -e "Downloading $model_name..."
    if command -v wget &> /dev/null; then
        wget --show-progress -qO "$output_file" "$download_url"
    elif command -v curl &> /dev/null; then
        curl -L --progress-bar -o "$output_file" "$download_url"
    else
        echo -e "${RED}No downloader found.${NC}"
        return 1
    fi

    if [[ $? -eq 0 ]]; then
        echo -e "${GREEN}Successfully downloaded $model_name${NC}"
    else
        echo -e "${RED}Failed to download $model_name${NC}"
    fi
}

# ---------- Install Recommended Model (dvine) ----------
install_dvine_model() {
    echo -e "${BLUE}Installing recommended model: ${MODEL_NAME}...${NC}"
    log "INFO" "Starting dvine model installation"
    download_civitai_model "$MODEL_NAME" "$MODEL_VERSION_ID"
}

# ---------- Install Multiple Models from BoRnNo0b ----------
install_bornnoob_models() {
    echo -e "${BLUE}Select models to install from BoRnNo0b's collection:${NC}"

    local model_names=()
    for name in "${!BORNNOOB_MODELS[@]}"; do
        model_names+=("$name")
    done

    # Let user choose multiple models (or all)
    local selected
    selected=$(gum choose --no-limit --header "Use space to select, enter to confirm" "${model_names[@]}")

    if [[ -z "$selected" ]]; then
        echo -e "${YELLOW}No models selected.${NC}"
        return
    fi

    # Process each selected model
    echo "$selected" | while read -r model; do
        local version_id="${BORNNOOB_MODELS[$model]}"
        if [[ "$version_id" == "CHANGE_ME" ]]; then
            echo -e "${RED}Version ID for '$model' is not set. Please edit the script and add the correct version ID from CivitAI.${NC}"
            log "WARNING" "Skipping $model: version ID not set"
            continue
        fi
        download_civitai_model "$model" "$version_id"
    done
}

# ---------- Full installation routines ----------
install_repo() {
    local repo_url="$1"
    local repo_dir="$2"
    local req_file="$3"
    local upgrade_flag="${4:-}"

    clone_repo "$repo_url" "$repo_dir"
    create_venv "$repo_dir"
    # Configure webui-user.sh before installing PyTorch so that TORCH_COMMAND is available if ever needed
    configure_webui_user_sh "$repo_dir"
    install_pytorch "$repo_dir"
    install_requirements "$repo_dir" "$req_file" "$upgrade_flag"
    apply_best_tips "$repo_dir"
    echo -e "${GREEN}Installation complete for ${repo_dir}.${NC}"
    log "INFO" "Installation complete for $repo_dir"

    if [[ -z "${NON_INTERACTIVE:-}" ]]; then
        if gum confirm "Would you like to install the recommended 'dvine' model now?"; then
            install_dvine_model
        fi
    fi
}

install_comfy() {
    mkdir -p "$DIR_COMFY"
    clone_repo "$REPO_COMFY" "$DIR_COMFY"
    create_venv "$DIR_COMFY"
    install_pytorch "$DIR_COMFY"

    local manager_dir="${DIR_COMFY}/custom_nodes/ComfyUI-Manager"
    if [[ -d "$manager_dir" ]]; then
        verify_repo_url "$REPO_COMFY_MANAGER" "$manager_dir"
        echo -e "${YELLOW}ComfyUI-Manager already exists, pulling updates...${NC}"
        log "INFO" "Pulling ComfyUI-Manager updates"
        (cd "$manager_dir" && git pull)
    else
        echo -e "${GREEN}Cloning ComfyUI-Manager...${NC}"
        log "INFO" "Cloning ComfyUI-Manager"
        git clone "$REPO_COMFY_MANAGER" "$manager_dir"
    fi

    if [[ -f "${manager_dir}/requirements.txt" ]]; then
        echo -e "${GREEN}Installing ComfyUI-Manager dependencies...${NC}"
        log "INFO" "Installing ComfyUI-Manager requirements"
        (cd "$DIR_COMFY" && source "${VENV_DIR}/bin/activate" && pip install -r "${manager_dir}/requirements.txt")
    fi

    if [[ -f "$DIR_COMFY/extra_model_paths.yaml" ]]; then
        sed -i '/base_path:/c\ base_path: /usr/local/share/comfyui' "$DIR_COMFY/extra_model_paths.yaml" 2>/dev/null || true
    fi

    apply_comfy_tips "$DIR_COMFY"
    echo -e "${GREEN}ComfyUI installation complete.${NC}"
    log "INFO" "ComfyUI installation complete"

    if [[ -z "${NON_INTERACTIVE:-}" ]]; then
        if gum confirm "Would you like to install the recommended 'dvine' model now?"; then
            install_dvine_model
        fi
    fi
}

auto_configure() {
    local repo_dir="$1"
    if [[ ! -d "$repo_dir" ]]; then
        echo -e "${RED}Repository ${repo_dir} not found.${NC}"
        log "ERROR" "auto_configure: $repo_dir not found"
        return 1
    fi
    detect_gpu
    # Ensure webui-user.sh is properly configured for this repo
    if [[ "$repo_dir" != "$DIR_COMFY" ]]; then
        configure_webui_user_sh "$repo_dir"
    fi
    install_pytorch "$repo_dir"
    if [[ "$repo_dir" == "$DIR_COMFY" ]]; then
        apply_comfy_tips "$repo_dir"
        if [[ "$GPU_TYPE" == "nvidia" ]]; then
            echo -e "${GREEN}NVIDIA GPU with ${GPU_VRAM}MB VRAM${NC}"
        elif [[ "$GPU_TYPE" == "amd" ]]; then
            echo -e "${GREEN}AMD GPU with ${GPU_GFX} architecture${NC}"
        fi
    else
        apply_best_tips "$repo_dir"
    fi
    echo -e "${GREEN}Hardware configuration complete for ${repo_dir}.${NC}"
    log "INFO" "Hardware configuration complete for $repo_dir"
}

# ---------- Interactive Menu ----------
main_menu() {
    local temp=$(get_gpu_temp)
    local header="Stable Diffusion WebUI Manager v1.1.6 | Attention: $ATTENTION_MODE | GPU Temp: $temp"
    local choice
    choice=$(gum choose --header "$header" \
        "Install AUTOMATIC1111 WebUI" \
        "Install Forge WebUI" \
        "Install ComfyUI + Manager" \
        "Install BOTH A1111 & Forge (with model linking)" \
        "Link A1111/Forge models (symlinks)" \
        "Link models to ComfyUI (extra_model_paths.yaml)" \
        "Set up unified model hub (central storage)" \
        "Install recommended model (dvine)" \
        "Install multiple models from BoRnNo0b" \
        "Auto-optimize configs for installed UIs" \
        "Fix Paths (repair broken links/configs)" \
        "Toggle attention mechanism (xformers ↔ sdp-no-mem)" \
        "Hardware status dashboard" \
        "Start thermal monitor (background)" \
        "Stop thermal monitor" \
        "Check for dependency conflicts" \
        "Update all dependencies (pip upgrade)" \
        "Reset venv for AUTOMATIC1111" \
        "Reset venv for Forge" \
        "Reset venv for ComfyUI" \
        "Install/Reinstall PyTorch (choose GPU type)" \
        "Install system dependencies" \
        "Exit")

    case "$choice" in
        "Install AUTOMATIC1111 WebUI")
            pre_install_checks
            detect_gpu
            install_repo "$REPO_A1111" "$DIR_A1111" "$REQUIREMENTS_FILE"
            if [[ -n "${NON_INTERACTIVE:-}" ]]; then
                pip cache purge || true
            fi
            ;;
        "Install Forge WebUI")
            pre_install_checks
            detect_gpu
            install_repo "$REPO_FORGE" "$DIR_FORGE" "$REQUIREMENTS_FILE"
            if [[ -n "${NON_INTERACTIVE:-}" ]]; then
                pip cache purge || true
            fi
            ;;
        "Install ComfyUI + Manager")
            pre_install_checks
            detect_gpu
            install_comfy
            if [[ -n "${NON_INTERACTIVE:-}" ]]; then
                pip cache purge || true
            fi
            if gum confirm "Do you want to link models from an existing A1111/Forge installation?"; then
                source_ui=$(gum choose "A1111" "Forge")
                generate_extra_model_paths "$DIR_COMFY" "$source_ui"
            fi
            ;;
        "Install BOTH A1111 & Forge (with model linking)")
            pre_install_checks
            detect_gpu
            install_repo "$REPO_A1111" "$DIR_A1111" "$REQUIREMENTS_FILE"
            install_repo "$REPO_FORGE" "$DIR_FORGE" "$REQUIREMENTS_FILE"
            master=$(gum choose --header "Choose the master directory (where models will be physically stored)" "AUTOMATIC1111" "Forge")
            if [[ "$master" == "AUTOMATIC1111" ]]; then
                link_models "$DIR_A1111" "$DIR_FORGE"
            else
                link_models "$DIR_FORGE" "$DIR_A1111"
            fi
            gum toast --title "Success" "Both WebUIs installed and models linked!"
            if [[ -n "${NON_INTERACTIVE:-}" ]]; then
                pip cache purge || true
            fi
            ;;
        "Link A1111/Forge models (symlinks)")
            master=$(gum choose --header "Choose the MASTER directory (where your models are actually stored)" "AUTOMATIC1111" "Forge")
            if [[ "$master" == "AUTOMATIC1111" ]]; then
                link_models "$DIR_A1111" "$DIR_FORGE"
            else
                link_models "$DIR_FORGE" "$DIR_A1111"
            fi
            gum toast --status success "Models linked successfully!"
            ;;
        "Link models to ComfyUI (extra_model_paths.yaml)")
            if [[ ! -d "$DIR_COMFY" ]]; then
                gum style --foreground 196 "ComfyUI not installed. Please install it first."
            else
                source_ui=$(gum choose "A1111" "Forge")
                generate_extra_model_paths "$DIR_COMFY" "$source_ui"
            fi
            ;;
        "Set up unified model hub (central storage)")
            setup_central_models
            gum toast --status success "Unified model hub configured!"
            ;;
        "Install recommended model (dvine)")
            install_dvine_model
            gum toast --status success "Model installation attempted!"
            ;;
        "Install multiple models from BoRnNo0b")
            install_bornnoob_models
            gum toast --status success "Batch installation attempted!"
            ;;
        "Auto-optimize configs for installed UIs")
            detect_gpu
            for dir in "$DIR_A1111" "$DIR_FORGE" "$DIR_COMFY"; do
                if [[ -d "$dir" ]]; then
                    if [[ "$dir" == "$DIR_COMFY" ]]; then
                        apply_comfy_tips "$dir"
                    else
                        # Ensure webui-user.sh has the right python and torch command
                        configure_webui_user_sh "$dir"
                        apply_best_tips "$dir"
                    fi
                fi
            done
            gum toast --status success "Optimal settings applied based on your $GPU_TYPE GPU!"
            ;;
        "Fix Paths (repair broken links/configs)")
            fix_paths
            gum toast --status success "Paths fixed!"
            ;;
        "Toggle attention mechanism (xformers ↔ sdp-no-mem)")
            toggle_attention
            ;;
        "Hardware status dashboard")
            show_hardware_dashboard
            gum confirm "Press Enter to continue"
            ;;
        "Start thermal monitor (background)")
            start_thermal_monitor
            ;;
        "Stop thermal monitor")
            stop_thermal_monitor
            ;;
        "Check for dependency conflicts")
            check_dependency_conflicts
            gum confirm "Press Enter to continue"
            ;;
        "Update all dependencies (pip upgrade)")
            update_all_deps
            gum toast --status success "Dependencies updated!"
            ;;
        "Reset venv for AUTOMATIC1111")
            reset_venv "$DIR_A1111"
            # After reset, reconfigure and reinstall PyTorch? The user can do it manually, but we can prompt
            if gum confirm "Do you want to reinstall PyTorch now?"; then
                install_pytorch "$DIR_A1111"
            fi
            ;;
        "Reset venv for Forge")
            reset_venv "$DIR_FORGE"
            if gum confirm "Do you want to reinstall PyTorch now?"; then
                install_pytorch "$DIR_FORGE"
            fi
            ;;
        "Reset venv for ComfyUI")
            reset_venv "$DIR_COMFY"
            if gum confirm "Do you want to reinstall PyTorch now?"; then
                install_pytorch "$DIR_COMFY"
            fi
            ;;
        "Install/Reinstall PyTorch (choose GPU type)")
            dir_choice=$(gum choose "AUTOMATIC1111" "Forge" "ComfyUI")
            case "$dir_choice" in
                "AUTOMATIC1111") target="$DIR_A1111" ;;
                "Forge") target="$DIR_FORGE" ;;
                "ComfyUI") target="$DIR_COMFY" ;;
            esac

            gpu_choice=$(gum choose "Auto-detect (recommended)" "NVIDIA (CUDA)" "AMD (ROCm)" "Intel (IPEX)" "CPU only")
            case "$gpu_choice" in
                "Auto-detect (recommended)")
                    detect_gpu
                    install_pytorch "$target"
                    ;;
                "NVIDIA (CUDA)")
                    install_pytorch "$target" "nvidia"
                    ;;
                "AMD (ROCm)")
                    install_pytorch "$target" "amd"
                    ;;
                "Intel (IPEX)")
                    install_pytorch "$target" "intel"
                    ;;
                "CPU only")
                    install_pytorch "$target" "cpu"
                    ;;
            esac
            ;;
        "Install system dependencies")
            pre_install_checks
            install_system_deps
            ;;
        "Exit")
            stop_thermal_monitor
            echo "Goodbye!"
            log "INFO" "Script terminated by user"
            exit 0
            ;;
        *)
            echo -e "${RED}Invalid option.${NC}"
            ;;
    esac
}

# ---------- Non-Interactive Mode ----------
handle_non_interactive() {
    local upgrade_flag="--upgrade"
    pre_install_checks
    case "${NON_INTERACTIVE:-}" in
        "a1111")
            detect_gpu
            install_repo "$REPO_A1111" "$DIR_A1111" "$REQUIREMENTS_FILE" "$upgrade_flag"
            ;;
        "forge")
            detect_gpu
            install_repo "$REPO_FORGE" "$DIR_FORGE" "$REQUIREMENTS_FILE" "$upgrade_flag"
            ;;
        "comfy")
            detect_gpu
            install_comfy
            if [[ -d "$DIR_A1111" ]]; then
                generate_extra_model_paths "$DIR_COMFY" "A1111"
            elif [[ -d "$DIR_FORGE" ]]; then
                generate_extra_model_paths "$DIR_COMFY" "Forge"
            fi
            ;;
        "both")
            detect_gpu
            install_repo "$REPO_A1111" "$DIR_A1111" "$REQUIREMENTS_FILE" "$upgrade_flag"
            install_repo "$REPO_FORGE" "$DIR_FORGE" "$REQUIREMENTS_FILE" "$upgrade_flag"
            link_models "$DIR_A1111" "$DIR_FORGE"
            ;;
        "dvine")
            install_dvine_model
            ;;
        "bornnoob")
            for name in "${!BORNNOOB_MODELS[@]}"; do
                local version_id="${BORNNOOB_MODELS[$name]}"
                if [[ "$version_id" != "CHANGE_ME" ]]; then
                    download_civitai_model "$name" "$version_id"
                fi
            done
            ;;
        "configure")
            detect_gpu
            for dir in "$DIR_A1111" "$DIR_FORGE" "$DIR_COMFY"; do
                if [[ -d "$dir" ]]; then
                    auto_configure "$dir"
                fi
            done
            ;;
    esac

    pip cache purge || true

    if [[ ! -d "$DIR_A1111" && ! -d "$DIR_FORGE" ]]; then
        echo "Cleaning up unused repositories..."
        if [[ -d "$REPO_A1111" ]]; then
            rm -rf "$REPO_A1111"
        fi
        if [[ -d "$REPO_FORGE" ]]; then
            rm -rf "$REPO_FORGE"
        fi
    fi

    if [[ ! -d "$DIR_COMFY" ]]; then
        if [[ -d "$REPO_COMFY_MANAGER" ]]; then
            rm -rf "$REPO_COMFY_MANAGER"
        fi
    fi
}

# ---------- Initialization ----------
check_gum
detect_python

# ---------- Main Execution ----------
if [[ -n "${NON_INTERACTIVE:-}" ]]; then
    handle_non_interactive
else
    while true; do
        main_menu
        echo
        gum confirm "Return to main menu?" --affirmative="Yes" --negative="Exit" || break
    done
fi

log "INFO" "Script finished successfully"