#!/usr/bin/env bash
# Desktop environment setup for ComfyUI + TRELLIS + UniRig + Blender
# Target: Arch Linux with NVIDIA RTX 4070 (CUDA 12.x)
#
# Prerequisites:
#   - NVIDIA drivers installed (nvidia, nvidia-utils)
#   - CUDA toolkit installed (cuda, cudnn)
#   - uv installed (https://astral.sh/uv)
#
# Usage: ./scripts/setup.sh [--comfyui-only | --unirig-only | --all]
set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_DIR="$(dirname "$SCRIPT_DIR")"
COMFYUI_DIR="$PROJECT_DIR/ComfyUI"

RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
NC='\033[0m'

info()  { echo -e "${GREEN}[INFO]${NC} $*"; }
warn()  { echo -e "${YELLOW}[WARN]${NC} $*"; }
error() { echo -e "${RED}[ERROR]${NC} $*" >&2; }

check_prerequisites() {
    info "Checking prerequisites..."

    if ! command -v nvidia-smi &>/dev/null; then
        error "nvidia-smi not found. Install NVIDIA drivers: sudo pacman -S nvidia nvidia-utils"
        exit 1
    fi
    info "NVIDIA driver: $(nvidia-smi --query-gpu=driver_version --format=csv,noheader | head -1)"

    if ! command -v nvcc &>/dev/null; then
        error "nvcc not found. Install CUDA: sudo pacman -S cuda cudnn"
        exit 1
    fi
    info "CUDA: $(nvcc --version | grep release | awk '{print $6}')"

    if ! command -v uv &>/dev/null; then
        error "uv not found. Install: curl -LsSf https://astral.sh/uv/install.sh | sh"
        exit 1
    fi
    info "uv: $(uv --version)"

    if ! command -v git &>/dev/null; then
        error "git not found. Install: sudo pacman -S git"
        exit 1
    fi
}

install_comfyui() {
    info "Installing ComfyUI..."

    if [[ -d "$COMFYUI_DIR" ]]; then
        info "ComfyUI directory exists, pulling latest..."
        git -C "$COMFYUI_DIR" pull
    else
        git clone https://github.com/comfyanonymous/ComfyUI.git "$COMFYUI_DIR"
    fi

    cd "$COMFYUI_DIR"

    if [[ ! -d ".venv" ]]; then
        uv venv --python 3.11
    fi

    # shellcheck disable=SC1091
    source .venv/bin/activate
    uv pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu124
    uv pip install -r requirements.txt

    info "ComfyUI installed at $COMFYUI_DIR"
}

install_3d_pack() {
    info "Installing ComfyUI-3D-Pack (includes TRELLIS nodes)..."

    local nodes_dir="$COMFYUI_DIR/custom_nodes"
    mkdir -p "$nodes_dir"

    if [[ -d "$nodes_dir/ComfyUI-3D-Pack" ]]; then
        info "ComfyUI-3D-Pack exists, pulling latest..."
        git -C "$nodes_dir/ComfyUI-3D-Pack" pull
    else
        git clone https://github.com/MrForExample/ComfyUI-3D-Pack.git "$nodes_dir/ComfyUI-3D-Pack"
    fi

    cd "$nodes_dir/ComfyUI-3D-Pack"

    # Activate ComfyUI venv for shared deps
    # shellcheck disable=SC1091
    source "$COMFYUI_DIR/.venv/bin/activate"
    uv pip install -r requirements.txt

    if [[ -f "install.py" ]]; then
        python install.py
    fi

    info "ComfyUI-3D-Pack installed"
}

install_unirig() {
    info "Installing UniRig..."

    local unirig_dir="$PROJECT_DIR/UniRig"

    if [[ -d "$unirig_dir" ]]; then
        info "UniRig exists, pulling latest..."
        git -C "$unirig_dir" pull
    else
        git clone https://github.com/VAST-AI-Research/UniRig.git "$unirig_dir"
    fi

    cd "$unirig_dir"

    # UniRig uses its own venv (different torch/dep requirements possible)
    if [[ ! -d ".venv" ]]; then
        uv venv --python 3.11
    fi

    # shellcheck disable=SC1091
    source .venv/bin/activate
    uv pip install torch torchvision --index-url https://download.pytorch.org/whl/cu124
    uv pip install -r requirements.txt

    # spconv for sparse convolutions
    uv pip install spconv-cu124 || warn "spconv-cu124 install failed — may need manual build"

    # flash-attn for efficient attention
    uv pip install flash-attn --no-build-isolation || warn "flash-attn install failed — TRELLIS will fall back to xformers"

    info "UniRig installed at $unirig_dir"
}

install_blender() {
    info "Checking Blender..."

    if command -v blender &>/dev/null; then
        info "Blender: $(blender --version 2>/dev/null | head -1)"
    else
        warn "Blender not found. Install: sudo pacman -S blender"
        warn "Then install VRM Add-on: https://vrm-addon-for-blender.info/en/"
    fi
}

install_rclone() {
    info "Checking rclone..."

    if command -v rclone &>/dev/null; then
        info "rclone: $(rclone --version | head -1)"

        if rclone listremotes | grep -q "^gravenhollow:"; then
            info "rclone remote 'gravenhollow' already configured"
        else
            warn "rclone remote 'gravenhollow' not configured."
            warn "Run: rclone config create gravenhollow s3 provider=Other endpoint=https://gravenhollow.lab.daviestechlabs.io:30292 access_key_id=<key> secret_access_key=<secret>"
        fi
    else
        warn "rclone not found. Install: sudo pacman -S rclone"
    fi
}

install_ray() {
    info "Checking Ray..."

    # Install ray into ComfyUI venv
    # shellcheck disable=SC1091
    source "$COMFYUI_DIR/.venv/bin/activate"

    if python -c "import ray" 2>/dev/null; then
        info "Ray: $(python -c 'import ray; print(ray.__version__)')"
    else
        info "Installing Ray..."
        uv pip install "ray[default]>=2.53.0"
        info "Ray installed"
    fi
}

install_avatar_pipeline() {
    info "Installing avatar-pipeline package..."
    cd "$PROJECT_DIR"

    # shellcheck disable=SC1091
    source "$COMFYUI_DIR/.venv/bin/activate"
    uv pip install -e ".[dev]"

    info "avatar-pipeline installed (editable)"
}

# ── Main ──────────────────────────────────────────────────────────────

MODE="${1:---all}"

check_prerequisites

case "$MODE" in
    --comfyui-only)
        install_comfyui
        install_3d_pack
        ;;
    --unirig-only)
        install_unirig
        ;;
    --all)
        install_comfyui
        install_3d_pack
        install_unirig
        install_blender
        install_rclone
        install_ray
        install_avatar_pipeline
        ;;
    *)
        echo "Usage: $0 [--comfyui-only | --unirig-only | --all]"
        exit 1
        ;;
esac

info "Setup complete!"
info ""
info "Next steps:"
info "  1. Start ComfyUI:  cd $COMFYUI_DIR && source .venv/bin/activate && python main.py"
info "  2. Open browser:   http://localhost:8188"
info "  3. Build your image-to-VRM workflow in the ComfyUI node graph"
info "  4. Export workflow JSON to: $PROJECT_DIR/workflows/"
info "  5. Join Ray cluster: ./scripts/ray-join.sh"