#!/bin/bash
# CoderAI - OpenAI-compatible API server
# Copyright (C) 2026 Stefy Lanza <stefy@nexlab.net>
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program. If not, see <https://www.gnu.org/licenses/>.

# Build script for CoderAI - Supports NVIDIA (CUDA), Vulkan, OpenCL, and CPU backends
# Usage: ./build.sh [nvidia|vulkan|vulkan-nvidia|cuda|opencl|all] [--flash] [--venv <venv>] [--package]
# Default: all (installs all backends)
# --flash: Enable and install Flash Attention 2 (for NVIDIA GPUs)
# --venv <venv>: Specify custom virtual environment name

set -e

# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# Determine backend and flags
BACKEND="${1:-all}"
FLASH=false
CUSTOM_VENV=""
PACKAGE=false
DS4=false

# Parse arguments
i=1
for arg in "$@"; do
    case $arg in
        --flash)
            FLASH=true
            ;;
        --venv)
            i=$((i + 1))
            eval "CUSTOM_VENV=\${$i}"
            ;;
        --package)
            PACKAGE=true
            ;;
        --ds4)
            DS4=true
            ;;
    esac
    i=$((i + 1))
done

BACKEND=$(echo "$BACKEND" | tr '[:upper:]' '[:lower:]')

if [[ "$BACKEND" != "nvidia" && "$BACKEND" != "vulkan" && "$BACKEND" != "vulkan-nvidia" && "$BACKEND" != "cuda" && "$BACKEND" != "opencl" && "$BACKEND" != "all" ]]; then
    echo -e "${RED}Error: Invalid backend '$BACKEND'${NC}"
    echo "Usage: ./build.sh [nvidia|vulkan|vulkan-nvidia|cuda|opencl|all] [--flash]"
    echo "  nvidia       - Use PyTorch with CUDA for NVIDIA GPUs"
    echo "  vulkan      - Use llama-cpp-python with Vulkan for AMD GPUs"
    echo "  vulkan-nvidia - Use llama-cpp-python with Vulkan for NVIDIA GPU only"
    echo "  cuda        - Use llama-cpp-python with CUDA for NVIDIA GPUs"
    echo "  opencl      - Use stable-diffusion-cpp-python with OpenCL"
    echo "  all         - Install all backends (nvidia, cuda, vulkan, opencl, cpu) - DEFAULT"
    echo ""
    echo "Options:"
    echo "  --flash     - Install Flash Attention 2 for faster inference (NVIDIA only)"
    echo "  --ds4       - Clone + build the ds4 (DeepSeek V4) native engine"
    exit 1
fi

echo -e "${BLUE}========================================${NC}"
echo -e "${BLUE}  CoderAI Build Script${NC}"
echo -e "${BLUE}  Backend: ${GREEN}$BACKEND${NC}"
if [ "$FLASH" = true ]; then
    echo -e "${BLUE}  Flash Attention 2: ${GREEN}ENABLED${NC}"
fi
echo -e "${BLUE}========================================${NC}"
echo ""

# Check Python version
PYTHON_VERSION=$(python3 --version 2>&1 | grep -oP '\d+\.\d+' | head -1)
REQUIRED_VERSION="3.8"

if [ "$(printf '%s\n' "$REQUIRED_VERSION" "$PYTHON_VERSION" | sort -V | head -n1)" != "$REQUIRED_VERSION" ]; then
    echo -e "${RED}Error: Python 3.8+ required, found $PYTHON_VERSION${NC}"
    exit 1
fi

echo -e "${GREEN}✓ Python version: $PYTHON_VERSION${NC}"

# Determine cmake args for stable-diffusion-cpp-python.
# The pip release is missing the libwebm/build/ cmake submodule files.
# If libwebm-dev is installed system-wide we can link against it; otherwise disable WebM.
if ldconfig -p 2>/dev/null | grep -q "libwebm" || pkg-config --exists libwebm 2>/dev/null; then
    SD_CMAKE_ARGS="-DSD_USE_SYSTEM_WEBM=ON"
    echo -e "${GREEN}✓ Found system libwebm — stable-diffusion-cpp-python will use it${NC}"
else
    SD_CMAKE_ARGS="-DSD_WEBM=OFF"
    echo -e "${YELLOW}Note: libwebm-dev not found — WebM video output disabled for stable-diffusion-cpp-python${NC}"
    echo -e "${YELLOW}      Install libwebm-dev to enable WebM support${NC}"
fi

# Determine venv directory based on backend
if [ -n "$CUSTOM_VENV" ]; then
    VENV_DIR="$CUSTOM_VENV"
    echo -e "${BLUE}Using custom virtual environment: $VENV_DIR${NC}"
elif [ "$BACKEND" = "nvidia" ]; then
    VENV_DIR="venv_nvidia"
elif [ "$BACKEND" = "vulkan" ]; then
    VENV_DIR="venv_vulkan"
elif [ "$BACKEND" = "vulkan-nvidia" ]; then
    VENV_DIR="venv_vulkan_nvidia"
elif [ "$BACKEND" = "cuda" ]; then
    VENV_DIR="venv_cuda"
elif [ "$BACKEND" = "opencl" ]; then
    VENV_DIR="venv_opencl"
elif [ "$BACKEND" = "all" ]; then
    VENV_DIR="venv_all"
fi

# Create virtual environment if it doesn't exist
if [ -n "$CUSTOM_VENV" ]; then
    echo -e "${YELLOW}Creating custom virtual environment: $VENV_DIR${NC}"
else
    echo -e "${YELLOW}Creating virtual environment: $VENV_DIR${NC}"
fi
if [ ! -d "$VENV_DIR" ]; then
    python3 -m venv "$VENV_DIR"
    echo -e "${GREEN}✓ Created virtual environment: $VENV_DIR${NC}"
else
    echo -e "${YELLOW}Using existing virtual environment: $VENV_DIR${NC}"
fi

# Activate virtual environment
echo -e "${YELLOW}Activating virtual environment...${NC}"
source "$VENV_DIR/bin/activate"

# Force pip to use this venv and install packages
export PIP_NO_INPUT=1
export PIP_REQUIRE_VIRTUALENV=1

# Upgrade pip
echo -e "${YELLOW}Upgrading pip...${NC}"
pip install --upgrade pip

echo ""
echo -e "${BLUE}Installing dependencies for $BACKEND backend...${NC}"
echo ""

if [ "$BACKEND" = "nvidia" ]; then
    # NVIDIA/CUDA backend
    echo -e "${YELLOW}Installing PyTorch with CUDA support...${NC}"
    pip install "torch>=2.0.0" "torchvision>=0.15.0" "torchaudio>=2.0.0"
    
    echo -e "${YELLOW}Installing NVIDIA-specific requirements...${NC}"
    pip install -r requirements-nvidia.txt || {
        echo -e "${YELLOW}Warning: Some NVIDIA packages failed to install${NC}"
    }

    # Extended modality dependencies (video/audio/embeddings/upscaling)
    echo -e "${YELLOW}Installing extended modality dependencies...${NC}"
    pip install "imageio[ffmpeg]" scipy soundfile sentence-transformers \
        openai-whisper argostranslate edge-tts kokoro-tts timm || {
        echo -e "${YELLOW}Warning: Some optional modality packages failed${NC}"
    }
    pip install realesrgan basicsr || {
        echo -e "${YELLOW}Warning: realesrgan/basicsr failed (image upscaling optional)${NC}"
    }
    pip install audiocraft 2>/dev/null || {
        echo -e "${YELLOW}Note: audiocraft not installed (audio generation with MusicGen optional)${NC}"
    }

    # Optional quantization backends for diffusers image/video pipelines:
    #   optimum-quanto -> enables 2-bit (int2) per-component quantization
    #   gguf           -> enables loading GGUF-quantized components (Q5_K/Q6_K, etc.)
    # bitsandbytes (4-bit/8-bit) comes via requirements-nvidia.txt; these add the
    # extra widths that bitsandbytes cannot do.
    echo -e "${YELLOW}Installing optional quantization backends (2-bit / GGUF)...${NC}"
    pip install optimum-quanto gguf || {
        echo -e "${YELLOW}Note: optimum-quanto/gguf not installed (2-bit and GGUF 5/6-bit quantization optional)${NC}"
    }

    # Install Flash Attention 2 if requested
    if [ "$FLASH" = true ]; then
        echo ""
        echo -e "${YELLOW}Installing Flash Attention 2...${NC}"
        echo -e "${YELLOW}This may take several minutes and requires CUDA 11.6+${NC}"
        MAX_JOBS=5 NVCC_THREADS=2  pip install flash-attn --no-build-isolation || {
            echo -e "${RED}Warning: Flash Attention 2 installation failed${NC}"
            echo -e "${YELLOW}Requirements: CUDA 11.6+, Linux, Ampere/Ada/Hopper GPU${NC}"
            echo -e "${YELLOW}Continuing without Flash Attention...${NC}"
        }
    fi
    
    echo ""
    echo -e "${GREEN}========================================${NC}"
    echo -e "${GREEN}  NVIDIA/CUDA build complete!${NC}"
    echo -e "${GREEN}========================================${NC}"
    echo ""
    echo "Usage:"
    echo "  source $VENV_DIR/bin/activate"
    echo "  python coderai --model <huggingface-model-name>"
    if [ "$FLASH" = true ]; then
        echo ""
        echo "Flash Attention 2 enabled - use --flash-attn flag when running"
    fi
    echo ""
    echo "Example:"
    echo "  python coderai --model microsoft/DialoGPT-medium"
    echo ""
    echo -e "${GREEN}========================================${NC}"
    echo -e "${GREEN}  NVIDIA/CUDA build complete!${NC}"
    echo -e "${GREEN}========================================${NC}"
    echo ""
    echo "Usage:"
    echo "  source $VENV_DIR/bin/activate"
    echo "  python coderai --model <huggingface-model-name>"
    echo ""
    echo "Example:"
    echo "  python coderai --model microsoft/DialoGPT-medium"
    echo ""
    
elif [ "$BACKEND" = "vulkan" ]; then
    # Vulkan backend (all GPUs)
    echo -e "${YELLOW}Installing llama-cpp-python with Vulkan support (all GPUs)...${NC}"
    
    # Check for required Vulkan development libraries
    if ! pkg-config --exists vulkan 2>/dev/null; then
        echo -e "${YELLOW}Warning: Vulkan development libraries not found via pkg-config${NC}"
        echo -e "${YELLOW}You may need to install Vulkan drivers and SDK:${NC}"
        echo "  Debian/Ubuntu: sudo apt install libvulkan-dev vulkan-tools"
        echo "  Fedora: sudo dnf install vulkan-loader-devel vulkan-tools"
        echo "  Arch: sudo pacman -S vulkan-headers vulkan-icd-loader"
        echo ""
        echo -e "${YELLOW}Attempting installation anyway...${NC}"
    fi
    
    # Check for glslc (Vulkan shader compiler)
    GLSLC_CMD=""
    if command -v glslc &> /dev/null; then
        GLSLC_CMD="glslc"
    elif command -v glslangValidator &> /dev/null; then
        GLSLC_CMD="glslangValidator"
    fi
    
    if [ -z "$GLSLC_CMD" ]; then
        echo -e "${YELLOW}Warning: glslc/glslangValidator not found in PATH${NC}"
    else
        echo -e "${GREEN}✓ Found Vulkan shader compiler: $GLSLC_CMD${NC}"
    fi
    
    # Build with Vulkan support (add CUDA too if available)
    echo -e "${YELLOW}Building llama-cpp-python with Vulkan support...${NC}"
    _LLAMA_CMAKE="-DGGML_VULKAN=ON"
    if command -v nvcc &> /dev/null || [ -d "/usr/local/cuda" ]; then
        _LLAMA_CMAKE="$_LLAMA_CMAKE -DGGML_CUDA=ON"
        echo -e "${GREEN}  ✓ Also enabling CUDA support (NVIDIA detected)${NC}"
    fi
    CMAKE_ARGS="$_LLAMA_CMAKE" pip install --upgrade llama-cpp-python --no-cache-dir || {
        echo -e "${RED}Build failed!${NC}"
        exit 1
    }
    
    echo -e "${YELLOW}Installing Vulkan-specific requirements...${NC}"
    pip install -r requirements-vulkan.txt
    echo -e "${YELLOW}Building whispercpp with Vulkan support for GPU-accelerated transcription...${NC}"
    
    # First, uninstall any existing whispercpp (pip version doesn't have Vulkan)
    pip uninstall -y whispercpp 2>/dev/null || true
    
    # Clone and build whisper.cpp with Vulkan for Python bindings
    WHISPERCPP_DIR="$HOME/whisper.cpp"
    if [ ! -d "$WHISPERCPP_DIR" ]; then
        echo "Cloning whisper.cpp..."
        git clone --depth 1 https://github.com/ggerganov/whisper.cpp "$WHISPERCPP_DIR" 2>/dev/null || {
            echo -e "${YELLOW}Warning: Could not clone whisper.cpp${NC}"
        }
    fi
    
    if [ -d "$WHISPERCPP_DIR/bindings/python" ]; then
        cd "$WHISPERCPP_DIR/bindings/python"
        
        # Build with Vulkan support
        # Set CMAKE_ARGS to enable Vulkan for ggml (whisper uses ggml library internally)
        CMAKE_ARGS="-DWHISPER_VULKAN=ON -DGGML_VULKAN=ON" pip install . --no-cache-dir --force-reinstall 2>/dev/null || {
            # If Vulkan build fails, try without (will fall back to CPU)
            echo -e "${YELLOW}Warning: whispercpp Vulkan build failed, will use CPU${NC}"
            pip install . --no-cache-dir --force-reinstall 2>/dev/null || {
                echo -e "${YELLOW}Warning: Could not install whispercpp at all${NC}"
            }
        }
        cd "$OLDPWD"
        echo -e "${GREEN}✓ whispercpp with Vulkan support installed!${NC}"
    else
        echo -e "${YELLOW}Warning: whisper.cpp Python bindings not found${NC}"
    fi
    
    # Also build the main whisper.cpp C++ with Vulkan for standalone usage
    echo -e "${YELLOW}Building whisper.cpp C++ with Vulkan support (optional)...${NC}"
    WHISPER_DIR="$HOME/whisper.cpp"
    if [ -d "$WHISPER_DIR" ]; then
        echo "Using existing whisper.cpp installation"
    else
        echo "Cloning whisper.cpp..."
        git clone https://github.com/ggerganov/whisper.cpp "$WHISPER_DIR" 2>/dev/null || {
            echo -e "${YELLOW}Warning: Could not clone whisper.cpp. Audio transcription will use CPU.${NC}"
        }
    fi
    
    if [ -d "$WHISPER_DIR" ]; then
        cd "$WHISPER_DIR"
        mkdir -p build 2>/dev/null
        cd build
        cmake -DGGML_VULKAN=ON .. >/dev/null 2>&1 || {
            echo -e "${YELLOW}Warning: Vulkan build failed, building with OpenBLAS${NC}"
            cmake -DBUILD_SHARED_LIBS=ON .. >/dev/null 2>&1
        }
        make -j$(nproc) >/dev/null 2>&1 || {
            echo -e "${YELLOW}Warning: Build failed. Audio transcription will use CPU.${NC}"
        }
        cd "$OLDPWD"
        
        if [ ! -f "$WHISPER_DIR/models/ggml-base.bin" ]; then
            echo "Downloading Whisper base model..."
            bash "$WHISPER_DIR/models/download-ggml-model.sh" base 2>/dev/null || {
                echo -e "${YELLOW}Warning: Could not download Whisper model.${NC}"
            }
        fi
        echo -e "${GREEN}✓ whisper.cpp ready for audio transcription!${NC}"
    fi
    
    echo ""
    echo -e "${GREEN}========================================${NC}"
    echo -e "${GREEN}  Vulkan build complete!${NC}"
    echo -e "${GREEN}========================================${NC}"
    echo ""
    echo "Usage:"
    echo "  python coderai --model <gguf-model> --backend vulkan"
    echo ""
    
elif [ "$BACKEND" = "vulkan-nvidia" ]; then
    # Vulkan backend (NVIDIA only)
    echo -e "${YELLOW}Installing llama-cpp-python with Vulkan support (NVIDIA-only)...${NC}"
    
    # Check for required Vulkan development libraries
    if ! pkg-config --exists vulkan 2>/dev/null; then
        echo -e "${YELLOW}Warning: Vulkan development libraries not found via pkg-config${NC}"
    fi
    
    # Check for glslc (Vulkan shader compiler)
    GLSLC_CMD=""
    if command -v glslc &> /dev/null; then
        GLSLC_CMD="glslc"
    elif command -v glslangValidator &> /dev/null; then
        GLSLC_CMD="glslangValidator"
    fi
    
    if [ -z "$GLSLC_CMD" ]; then
        echo -e "${YELLOW}Warning: glslc/glslangValidator not found in PATH${NC}"
    else
        echo -e "${GREEN}✓ Found Vulkan shader compiler: $GLSLC_CMD${NC}"
    fi
    
    # Build with Vulkan support (add CUDA too if available)
    # Note: llama.cpp doesn't have a compile-time option to disable specific GPUs
    # The device selection happens at runtime via environment variables
    echo -e "${YELLOW}Building llama-cpp-python with Vulkan support...${NC}"
    _LLAMA_CMAKE="-DGGML_VULKAN=ON"
    if command -v nvcc &> /dev/null || [ -d "/usr/local/cuda" ]; then
        _LLAMA_CMAKE="$_LLAMA_CMAKE -DGGML_CUDA=ON"
        echo -e "${GREEN}  ✓ Also enabling CUDA support (NVIDIA detected)${NC}"
    fi
    CMAKE_ARGS="$_LLAMA_CMAKE" pip install --upgrade llama-cpp-python --no-cache-dir || {
        echo -e "${RED}Build failed!${NC}"
        exit 1
    }
    
    echo -e "${YELLOW}Installing Vulkan-specific requirements...${NC}"
    pip install -r requirements-vulkan.txt
    
    # Build whispercpp Python package with Vulkan support for GPU-accelerated audio transcription
    echo -e "${YELLOW}Building whispercpp with Vulkan support for GPU-accelerated transcription...${NC}"
    pip uninstall -y whispercpp 2>/dev/null || true
    WHISPERCPP_DIR="$HOME/whisper.cpp"
    if [ ! -d "$WHISPERCPP_DIR" ]; then
        git clone --depth 1 https://github.com/ggerganov/whisper.cpp "$WHISPERCPP_DIR" 2>/dev/null || true
    fi
    if [ -d "$WHISPERCPP_DIR/bindings/python" ]; then
        cd "$WHISPERCPP_DIR/bindings/python"
        CMAKE_ARGS="-DWHISPER_VULKAN=ON -DGGML_VULKAN=ON" pip install . --no-cache-dir --force-reinstall 2>/dev/null || {
            pip install . --no-cache-dir --force-reinstall 2>/dev/null || true
        }
        cd "$OLDPWD"
        echo -e "${GREEN}✓ whispercpp with Vulkan support installed!${NC}"
    fi
    
    echo ""
    echo -e "${GREEN}========================================${NC}"
    echo -e "${GREEN}  Vulkan (NVIDIA-only) build complete!${NC}"
    echo -e "${GREEN}========================================${NC}"
    echo ""
    echo "Usage:"
    echo "  VK_ICD_FILENAMES=/usr/share/vulkan/icd.d/nvidia_icd.json \\"
    echo "  python coderai --model <gguf-model> --backend vulkan"
    echo ""
    echo "Note: This build includes both AMD and NVIDIA Vulkan support."
    echo "      At runtime, use VK_ICD_FILENAMES to select only NVIDIA."
    echo ""
 
elif [ "$BACKEND" = "cuda" ]; then
    # llama-cpp-python with CUDA backend (NVIDIA only)
    echo -e "${YELLOW}Installing llama-cpp-python with CUDA support...${NC}"
    
    # Check for CUDA toolkit
    if ! command -v nvcc &> /dev/null; then
        echo -e "${YELLOW}Warning: CUDA toolkit (nvcc) not found in PATH${NC}"
        echo -e "${YELLOW}You may need to install CUDA toolkit:${NC}"
        echo "  Download from: https://developer.nvidia.com/cuda-downloads"
    else
        CUDA_VERSION=$(nvcc --version | grep "release" | sed -n 's/.*release \([0-9.]*\),.*/\1/p')
        echo -e "${GREEN}✓ Found CUDA $CUDA_VERSION${NC}"
    fi
    
    # Check for CUDA libraries
    if [ -d "/usr/local/cuda" ]; then
        echo -e "${GREEN}✓ Found CUDA at /usr/local/cuda${NC}"
    fi
    
    # Build llama-cpp-python with CUDA support (add Vulkan too if available)
    echo -e "${YELLOW}Building llama-cpp-python with CUDA support...${NC}"
    echo -e "${YELLOW}This may take several minutes...${NC}"
    _LLAMA_CMAKE="-DGGML_CUDA=ON"
    if pkg-config --exists vulkan 2>/dev/null; then
        _LLAMA_CMAKE="$_LLAMA_CMAKE -DGGML_VULKAN=ON"
        echo -e "${GREEN}  ✓ Also enabling Vulkan support (Vulkan detected)${NC}"
    fi
    CMAKE_ARGS="$_LLAMA_CMAKE" pip install --upgrade llama-cpp-python --no-cache-dir || {
        echo ""
        echo -e "${RED}Build failed!${NC}"
        echo -e "${YELLOW}Make sure CUDA toolkit is installed:${NC}"
        echo "  sudo apt install cuda-toolkit-12"
        echo "  or"
        echo "  Download from: https://developer.nvidia.com/cuda-downloads"
        exit 1
    }
    
    echo -e "${YELLOW}Installing Vulkan-specific requirements...${NC}"
    pip install -r requirements-vulkan.txt
    
    echo ""
    echo -e "${GREEN}========================================${NC}"
    echo -e "${GREEN}  llama-cpp-python CUDA build complete!${NC}"
    echo -e "${GREEN}========================================${NC}"
    echo ""
    echo "Usage:"
    echo "  source $VENV_DIR/bin/activate"
    echo "  python coderai --model <gguf-model> --backend vulkan --vulkan-device 0"
    echo ""
    echo "Note: With CUDA backend, llama-cpp-python will only use NVIDIA GPUs."
    echo ""
elif [ "$BACKEND" = "opencl" ]; then
    # stable-diffusion-cpp-python with OpenCL backend
    echo -e "${YELLOW}Installing stable-diffusion-cpp-python with OpenCL support...${NC}"
    
    # Check for OpenCL
    if ! command -v clinfo &> /dev/null && ! ls /usr/lib/*/libOpenCL* &> /dev/null; then
        echo -e "${YELLOW}Warning: OpenCL not found in system${NC}"
        echo -e "${YELLOW}You may need to install OpenCL runtime:${NC}"
        echo "  Debian/Ubuntu: sudo apt install ocl-icd-opencl-dev"
        echo "  Fedora: sudo dnf install ocl-icd-devel"
    else
        echo -e "${GREEN}✓ Found OpenCL${NC}"
    fi
    
    # Install base requirements
    echo -e "${YELLOW}Installing base requirements...${NC}"
    pip install -r requirements.txt
    
    # Install stable-diffusion-cpp-python with OpenCL (disable WebM to avoid libwebm cmake issue)
    echo -e "${YELLOW}Installing stable-diffusion-cpp-python with OpenCL...${NC}"
    CMAKE_ARGS="$SD_CMAKE_ARGS" pip install stable-diffusion-cpp-python || {
        echo ""
        echo -e "${YELLOW}Note: If stable-diffusion-cpp-python is not available with pip,${NC}"
        echo -e "${YELLOW}you may need to build from source.${NC}"
    }
    
    # Install additional requirements for OpenCL
    echo -e "${YELLOW}Installing additional requirements for OpenCL...${NC}"
    pip install numpy pillow
    
    echo ""
    echo -e "${GREEN}========================================${NC}"
    echo -e "${GREEN}  OpenCL build complete!${NC}"
    echo -e "${GREEN}========================================${NC}"
    echo ""
    echo "Usage:"
    echo "  source $VENV_DIR/bin/activate"
    echo "  python coderai --model <model> --image-backend opencl"
    echo ""
    echo "Note: With OpenCL backend, stable-diffusion-cpp-python can use various GPUs."
    echo ""

elif [ "$BACKEND" = "all" ]; then
    # Install ALL backends: nvidia (CUDA), vulkan, opencl, and cpu
    echo -e "${BLUE}========================================${NC}"
    echo -e "${BLUE}  Installing ALL backends${NC}"
    echo -e "${BLUE}  (NVIDIA/CUDA, Vulkan, OpenCL, CPU)${NC}"
    echo -e "${BLUE}========================================${NC}"
    echo ""
    
    # Install base requirements
    echo -e "${YELLOW}Installing base requirements...${NC}"
    pip install --upgrade pip
    
    # Install requirements with error handling for problematic packages
    echo -e "${YELLOW}Installing core dependencies...${NC}"
    pip install -r requirements.txt || {
        echo -e "${YELLOW}Some packages failed to install, trying individually...${NC}"
        
        # Install core packages that should always work
        pip install fastapi uvicorn pydantic requests python-multipart psutil || {
            echo -e "${RED}Failed to install core dependencies${NC}"
            exit 1
        }
        
        # Try optional packages individually
        echo -e "${YELLOW}Installing optional packages...${NC}"
        pip install transformers accelerate diffusers safetensors || echo -e "${YELLOW}Warning: Some ML packages failed${NC}"
        pip install faster-whisper || echo -e "${YELLOW}Warning: faster-whisper failed${NC}"
        pip install whispercpp || echo -e "${YELLOW}Warning: whispercpp failed${NC}"
        pip install litellm || echo -e "${YELLOW}Warning: litellm failed${NC}"
        
        pip install setproctitle || echo -e "${YELLOW}Warning: setproctitle failed (optional)${NC}"

        # Try stable-diffusion-cpp-python (disable WebM to avoid missing libwebm cmake submodule)
        # Use CUDA if available (detected later in this block, check nvcc now)
        if command -v nvcc &> /dev/null || [ -d "/usr/local/cuda" ]; then
            CMAKE_ARGS="$SD_CMAKE_ARGS -DSD_CUDA=ON" pip install stable-diffusion-cpp-python --no-cache-dir || \
            CMAKE_ARGS="$SD_CMAKE_ARGS" pip install stable-diffusion-cpp-python || \
            echo -e "${YELLOW}Warning: stable-diffusion-cpp-python failed (optional)${NC}"
        else
            CMAKE_ARGS="$SD_CMAKE_ARGS" pip install stable-diffusion-cpp-python || echo -e "${YELLOW}Warning: stable-diffusion-cpp-python failed (optional)${NC}"
        fi
    }
    
    # Install PyTorch with CUDA support (for nvidia backend)
    echo -e "${YELLOW}Installing PyTorch with CUDA support (NVIDIA backend)...${NC}"
    pip install torch torchvision torchaudio || {
        echo -e "${YELLOW}Warning: PyTorch installation failed, will try CPU version${NC}"
        pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu || {
            echo -e "${RED}Failed to install PyTorch${NC}"
            exit 1
        }
    }
    
    echo -e "${YELLOW}Installing NVIDIA-specific requirements...${NC}"
    pip install -r requirements-nvidia.txt || {
        echo -e "${YELLOW}Warning: Some NVIDIA packages failed to install${NC}"
    }

    # Extended modality dependencies (video/audio/embeddings/upscaling)
    echo -e "${YELLOW}Installing extended modality dependencies...${NC}"
    pip install "imageio[ffmpeg]" scipy soundfile sentence-transformers \
        openai-whisper argostranslate edge-tts kokoro-tts timm || {
        echo -e "${YELLOW}Warning: Some optional modality packages failed${NC}"
    }
    pip install realesrgan basicsr || {
        echo -e "${YELLOW}Warning: realesrgan/basicsr failed (image upscaling optional)${NC}"
    }
    # audiocraft (MusicGen/AudioGen) — Meta package, may fail on some Python versions
    pip install audiocraft 2>/dev/null || {
        echo -e "${YELLOW}Note: audiocraft not installed (audio generation with MusicGen optional)${NC}"
        echo -e "${YELLOW}      Install manually: pip install audiocraft${NC}"
    }

    # Check for Vulkan development libraries
    VULKAN_AVAILABLE=false
    if pkg-config --exists vulkan 2>/dev/null; then
        VULKAN_AVAILABLE=true
        echo -e "${GREEN}✓ Found Vulkan development libraries${NC}"
    else
        echo -e "${YELLOW}Warning: Vulkan development libraries not found${NC}"
        echo -e "${YELLOW}  Vulkan support will be limited${NC}"
    fi
    
    # Check for CUDA
    CUDA_AVAILABLE=false
    if command -v nvcc &> /dev/null || [ -d "/usr/local/cuda" ]; then
        CUDA_AVAILABLE=true
        echo -e "${GREEN}✓ Found CUDA toolkit${NC}"
    else
        echo -e "${YELLOW}Warning: CUDA toolkit not found${NC}"
        echo -e "${YELLOW}  CUDA support will be limited${NC}"
    fi
    
    # Check for OpenCL
    OPENCL_AVAILABLE=false
    if command -v clinfo &> /dev/null || ls /usr/lib/*/libOpenCL* &> /dev/null 2>&1; then
        OPENCL_AVAILABLE=true
        echo -e "${GREEN}✓ Found OpenCL${NC}"
    else
        echo -e "${YELLOW}Warning: OpenCL not found${NC}"
        echo -e "${YELLOW}  OpenCL support will be limited${NC}"
    fi
    
    # Build llama-cpp-python with both CUDA and Vulkan support
    echo -e "${YELLOW}Building llama-cpp-python with CUDA and Vulkan support...${NC}"
    echo -e "${YELLOW}This may take several minutes...${NC}"
    
    # Determine CMAKE_ARGS based on available hardware
    CMAKE_ARGS=""
    if [ "$CUDA_AVAILABLE" = true ]; then
        CMAKE_ARGS="-DGGML_CUDA=ON"
        echo -e "${GREEN}  ✓ Enabling CUDA support${NC}"
    fi
    
    if [ "$VULKAN_AVAILABLE" = true ]; then
        if [ -n "$CMAKE_ARGS" ]; then
            CMAKE_ARGS="$CMAKE_ARGS -DGGML_VULKAN=ON"
        else
            CMAKE_ARGS="-DGGML_VULKAN=ON"
        fi
        echo -e "${GREEN}  ✓ Enabling Vulkan support${NC}"
    fi
    
    if [ -n "$CMAKE_ARGS" ]; then
        echo -e "${YELLOW}  Building with: $CMAKE_ARGS${NC}"
        CMAKE_ARGS="$CMAKE_ARGS" pip install --upgrade llama-cpp-python --no-cache-dir || {
            echo -e "${YELLOW}Warning: llama-cpp-python build failed, installing from pip${NC}"
            pip install llama-cpp-python
        }
    else
        echo -e "${YELLOW}Warning: No GPU backends available, installing CPU version${NC}"
        pip install llama-cpp-python
    fi
    
    # Install Vulkan-specific requirements
    echo -e "${YELLOW}Installing Vulkan-specific requirements...${NC}"
    pip install -r requirements-vulkan.txt || {
        echo -e "${YELLOW}Warning: Some Vulkan packages failed to install${NC}"
    }
    
    # Try to install stable-diffusion-cpp-python with CUDA+Vulkan (preferred) or fallbacks
    if [ "$CUDA_AVAILABLE" = true ] && [ "$VULKAN_AVAILABLE" = true ]; then
        echo -e "${YELLOW}Installing stable-diffusion-cpp-python with CUDA+Vulkan support...${NC}"
        CMAKE_ARGS="$SD_CMAKE_ARGS -DSD_CUDA=ON -DSD_VULKAN=ON" pip install stable-diffusion-cpp-python --no-cache-dir || {
            echo -e "${YELLOW}CUDA+Vulkan build failed, trying CUDA only...${NC}"
            CMAKE_ARGS="$SD_CMAKE_ARGS -DSD_CUDA=ON" pip install stable-diffusion-cpp-python --no-cache-dir || \
                echo -e "${YELLOW}Warning: stable-diffusion-cpp-python not available${NC}"
        }
    elif [ "$CUDA_AVAILABLE" = true ]; then
        echo -e "${YELLOW}Installing stable-diffusion-cpp-python with CUDA support...${NC}"
        CMAKE_ARGS="$SD_CMAKE_ARGS -DSD_CUDA=ON" pip install stable-diffusion-cpp-python --no-cache-dir || \
            echo -e "${YELLOW}Warning: stable-diffusion-cpp-python not available${NC}"
    elif [ "$VULKAN_AVAILABLE" = true ]; then
        echo -e "${YELLOW}Installing stable-diffusion-cpp-python with Vulkan support...${NC}"
        CMAKE_ARGS="$SD_CMAKE_ARGS -DSD_VULKAN=ON" pip install stable-diffusion-cpp-python --no-cache-dir || \
            echo -e "${YELLOW}Warning: stable-diffusion-cpp-python not available${NC}"
    elif [ "$OPENCL_AVAILABLE" = true ]; then
        echo -e "${YELLOW}Installing stable-diffusion-cpp-python with OpenCL support...${NC}"
        CMAKE_ARGS="$SD_CMAKE_ARGS -DSD_OPENCL=ON" pip install stable-diffusion-cpp-python --no-cache-dir || \
            echo -e "${YELLOW}Warning: stable-diffusion-cpp-python not available${NC}"
    else
        echo -e "${YELLOW}Skipping GPU-accelerated stable-diffusion-cpp-python - no GPU backend available${NC}"
    fi

    # Install additional requirements
    echo -e "${YELLOW}Installing additional requirements...${NC}"
    pip install numpy pillow || {
        echo -e "${YELLOW}Warning: Some additional packages failed${NC}"
    }

    # Install setproctitle for process naming (Python 3.13 compatible)
    echo -e "${YELLOW}Installing setproctitle...${NC}"
    pip install setproctitle || {
        echo -e "${YELLOW}Note: setproctitle failed to install (optional package, not critical)${NC}"
    }
    
    # Install Flash Attention 2 if requested and CUDA is available
    if [ "$FLASH" = true ] && [ "$CUDA_AVAILABLE" = true ]; then
        echo ""
        echo -e "${YELLOW}Installing Flash Attention 2...${NC}"
        echo -e "${YELLOW}This may take several minutes and requires CUDA 11.6+${NC}"
        MAX_JOBS=5 NVCC_THREADS=2 pip install flash-attn --no-build-isolation || {
            echo -e "${RED}Warning: Flash Attention 2 installation failed${NC}"
            echo -e "${YELLOW}Requirements: CUDA 11.6+, Linux, Ampere/Ada/Hopper GPU${NC}"
            echo -e "${YELLOW}Continuing without Flash Attention...${NC}"
        }
    elif [ "$FLASH" = true ]; then
        echo -e "${YELLOW}Warning: Flash Attention 2 requires CUDA backend${NC}"
        echo -e "${YELLOW}Skipping Flash Attention installation${NC}"
    fi
    
    echo ""
    echo -e "${GREEN}========================================${NC}"
    echo -e "${GREEN}  ALL backends build complete!${NC}"
    echo -e "${GREEN}========================================${NC}"
    echo ""
    echo "Available backends:"
    [ "$CUDA_AVAILABLE" = true ] && echo "  ✓ NVIDIA/CUDA (PyTorch)"
    [ "$CUDA_AVAILABLE" = true ] && echo "  ✓ CUDA (llama-cpp-python)"
    [ "$CUDA_AVAILABLE" = true ] && [ "$VULKAN_AVAILABLE" = true ] && echo "  ✓ CUDA+Vulkan (stable-diffusion-cpp-python)"
    [ "$CUDA_AVAILABLE" = true ] && [ "$VULKAN_AVAILABLE" != true ] && echo "  ✓ CUDA (stable-diffusion-cpp-python)"
    [ "$CUDA_AVAILABLE" != true ] && [ "$VULKAN_AVAILABLE" = true ] && echo "  ✓ Vulkan (stable-diffusion-cpp-python)"
    [ "$VULKAN_AVAILABLE" = true ] && echo "  ✓ Vulkan (llama-cpp-python)"
    [ "$OPENCL_AVAILABLE" = true ] && [ "$CUDA_AVAILABLE" != true ] && [ "$VULKAN_AVAILABLE" != true ] && echo "  ✓ OpenCL (stable-diffusion-cpp-python)"
    echo "  ✓ CPU (fallback for all)"
    if [ "$FLASH" = true ] && [ "$CUDA_AVAILABLE" = true ]; then
        echo ""
        echo "  ✓ Flash Attention 2 (NVIDIA)"
    fi
    echo ""
    echo "Usage:"
    echo "  source $VENV_DIR/bin/activate"
    echo ""
    echo "  # For text models with NVIDIA:"
    echo "  python coderai --model <model> --backend nvidia"
    if [ "$FLASH" = true ]; then
        echo "  python coderai --model <model> --backend nvidia --flash-attn"
    fi
    echo ""
    echo "  # For GGUF models with CUDA:"
    echo "  python coderai --model <gguf-model> --backend vulkan"
    echo ""
    echo "  # For GGUF models with Vulkan:"
    echo "  python coderai --model <gguf-model> --backend vulkan"
    echo ""
    echo "  # For image generation with OpenCL:"
    echo "  python coderai --model <model> --image-backend opencl"
    echo ""
fi

package_app() {
    echo -e "${YELLOW}Packaging CoderAI with PyInstaller...${NC}"
    pip install pyinstaller
    mkdir -p dist-package
    pyinstaller --clean --noconfirm --onefile --name coderai \
        --collect-all codai \
        --collect-all fastapi \
        --collect-all uvicorn \
        --collect-all pydantic \
        --collect-all transformers \
        --collect-all diffusers \
        --collect-all sentence_transformers \
        --collect-all whispercpp \
        --collect-all insightface \
        --collect-all onnxruntime \
        --collect-all PIL \
        coderai
    cp dist/coderai dist-package/coderai
    echo -e "${GREEN}✓ Packaged executable: dist-package/coderai${NC}"
    echo -e "${YELLOW}Note: The target machine must still provide compatible system GPU/runtime libraries.${NC}"
}

# Optionally clone + build ds4 (DeepSeek V4 native engine). Opt-in via --ds4.
# coderai can also auto-build this at runtime on first use, but doing it here lets
# the OCI/Docker packaging bundle the prebuilt ds4-server binary.
build_ds4() {
    local DS4_DIR="${CODERAI_DS4_DIR:-$HOME/.coderai/ds4}"
    echo -e "${YELLOW}Building ds4 (DeepSeek V4 engine) → $DS4_DIR ...${NC}"
    if [ ! -e "$DS4_DIR/Makefile" ]; then
        mkdir -p "$(dirname "$DS4_DIR")"
        git clone --depth 1 https://github.com/antirez/ds4 "$DS4_DIR" || {
            echo -e "${YELLOW}Warning: could not clone ds4; skipping.${NC}"; return 0; }
    fi
    local TARGET="cpu"
    if command -v nvcc &> /dev/null || [ -d "/usr/local/cuda" ]; then
        TARGET="cuda-generic"
    elif [ "$(uname -s)" = "Darwin" ]; then
        TARGET=""   # bare `make` builds the macOS Metal backend
    fi
    ( cd "$DS4_DIR" && make $TARGET ) || {
        echo -e "${YELLOW}Warning: ds4 build failed; it can still be built at runtime.${NC}"; return 0; }
    if [ -x "$DS4_DIR/ds4-server" ]; then
        echo -e "${GREEN}✓ ds4-server built at $DS4_DIR/ds4-server${NC}"
        echo -e "${YELLOW}Note: DeepSeek V4 weights are downloaded on first use (multi-GB).${NC}"
    fi
}

if [ "$DS4" = true ]; then
    build_ds4
fi

# Create .backend file to track which backend was used
echo "$BACKEND" > .backend

if [ "$PACKAGE" = true ]; then
    package_app
fi

echo -e "${GREEN}Build completed successfully!${NC}"
echo ""
echo "To activate the environment in the future, run:"
echo "  source $VENV_DIR/bin/activate"
