#!/bin/bash
# ==============================================================================
# WSL GPU Setup Script for Local LLM Agent
# ==============================================================================
# This script configures NVIDIA GPU support in WSL Ubuntu-24.04 for:
# - Docker GPU acceleration
# - vLLM backend
# - CUDA toolkit
#
# Prerequisites:
# - Windows 11 with WSL2
# - NVIDIA GPU with recent drivers (>= 525.xx)
# - Ubuntu-24.04 WSL distribution
#
# Usage:
#   ./scripts/setup-wsl-gpu.sh
#
# Author: ISEM Development Team
# Version: 1.0.0
# ==============================================================================

set -euo pipefail

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

# Logging functions
log_info() {
    echo -e "${BLUE}[INFO]${NC} $1"
}

log_success() {
    echo -e "${GREEN}[SUCCESS]${NC} $1"
}

log_warning() {
    echo -e "${YELLOW}[WARNING]${NC} $1"
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

# Check if running in WSL
check_wsl() {
    log_info "Checking WSL environment..."

    if [[ ! -f /proc/version ]] || ! grep -qi microsoft /proc/version; then
        log_error "This script must be run inside WSL (Windows Subsystem for Linux)"
        exit 1
    fi

    log_success "Running in WSL environment"
}

# Check Ubuntu version
check_ubuntu_version() {
    log_info "Checking Ubuntu version..."

    if [[ -f /etc/os-release ]]; then
        . /etc/os-release
        if [[ "$ID" == "ubuntu" ]]; then
            log_success "Ubuntu $VERSION_ID detected"
        else
            log_warning "Expected Ubuntu, found $ID. Proceeding anyway..."
        fi
    else
        log_warning "Could not detect OS version"
    fi
}

# Check for existing NVIDIA driver (should be provided by Windows)
check_nvidia_driver() {
    log_info "Checking NVIDIA driver..."

    if command -v nvidia-smi &> /dev/null; then
        DRIVER_VERSION=$(nvidia-smi --query-gpu=driver_version --format=csv,noheader | head -n1)
        log_success "NVIDIA driver version: $DRIVER_VERSION"

        # Show GPU info
        log_info "GPU Information:"
        nvidia-smi --query-gpu=name,memory.total --format=csv,noheader
    else
        log_error "NVIDIA driver not detected!"
        log_info "In WSL2, the NVIDIA driver is provided by Windows."
        log_info "Please ensure:"
        log_info "  1. NVIDIA Game Ready or Studio driver >= 525.xx is installed on Windows"
        log_info "  2. WSL is updated: wsl --update"
        log_info "  3. Restart WSL: wsl --shutdown"
        exit 1
    fi
}

# Install CUDA Toolkit
install_cuda_toolkit() {
    log_info "Installing CUDA Toolkit 12.6..."

    # Check if already installed
    if command -v nvcc &> /dev/null; then
        CUDA_VERSION=$(nvcc --version | grep release | awk '{print $6}' | cut -c2-)
        log_info "CUDA Toolkit $CUDA_VERSION already installed"

        read -p "Do you want to reinstall? (y/N) " -n 1 -r
        echo
        if [[ ! $REPLY =~ ^[Yy]$ ]]; then
            log_info "Skipping CUDA Toolkit installation"
            return
        fi
    fi

    # Add NVIDIA CUDA repository
    log_info "Adding NVIDIA CUDA repository..."

    # Download and install keyring
    wget -q https://developer.download.nvidia.com/compute/cuda/repos/wsl-ubuntu/x86_64/cuda-keyring_1.1-1_all.deb
    sudo dpkg -i cuda-keyring_1.1-1_all.deb
    rm cuda-keyring_1.1-1_all.deb

    # Update package list
    sudo apt-get update

    # Install CUDA Toolkit 12.6
    log_info "Installing CUDA Toolkit 12.6 (this may take a while)..."
    sudo apt-get install -y cuda-toolkit-12-6

    # Add CUDA to PATH
    if ! grep -q "cuda-12.6" ~/.bashrc; then
        echo '' >> ~/.bashrc
        echo '# CUDA 12.6' >> ~/.bashrc
        echo 'export PATH=/usr/local/cuda-12.6/bin:$PATH' >> ~/.bashrc
        echo 'export LD_LIBRARY_PATH=/usr/local/cuda-12.6/lib64:$LD_LIBRARY_PATH' >> ~/.bashrc
    fi

    # Export for current session
    export PATH=/usr/local/cuda-12.6/bin:$PATH
    export LD_LIBRARY_PATH=/usr/local/cuda-12.6/lib64:$LD_LIBRARY_PATH

    log_success "CUDA Toolkit 12.6 installed successfully"
}

# Install Docker (if not already installed)
install_docker() {
    log_info "Checking Docker installation..."

    if command -v docker &> /dev/null; then
        DOCKER_VERSION=$(docker --version | awk '{print $3}' | cut -d',' -f1)
        log_success "Docker $DOCKER_VERSION already installed"
    else
        log_info "Installing Docker..."

        # Install prerequisites
        sudo apt-get update
        sudo apt-get install -y ca-certificates curl gnupg

        # Add Docker's official GPG key
        sudo install -m 0755 -d /etc/apt/keyrings
        curl -fsSL https://download.docker.com/linux/ubuntu/gpg | sudo gpg --dearmor -o /etc/apt/keyrings/docker.gpg
        sudo chmod a+r /etc/apt/keyrings/docker.gpg

        # Add repository
        echo \
          "deb [arch="$(dpkg --print-architecture)" signed-by=/etc/apt/keyrings/docker.gpg] https://download.docker.com/linux/ubuntu \
          "$(. /etc/os-release && echo "$VERSION_CODENAME")" stable" | \
          sudo tee /etc/apt/sources.list.d/docker.list > /dev/null

        # Install Docker
        sudo apt-get update
        sudo apt-get install -y docker-ce docker-ce-cli containerd.io docker-buildx-plugin docker-compose-plugin

        # Add current user to docker group
        sudo usermod -aG docker $USER

        log_success "Docker installed successfully"
        log_warning "You may need to log out and log back in for docker group to take effect"
    fi
}

# Install NVIDIA Container Toolkit
install_nvidia_container_toolkit() {
    log_info "Installing NVIDIA Container Toolkit..."

    # Check if already installed
    if command -v nvidia-ctk &> /dev/null; then
        log_info "NVIDIA Container Toolkit already installed"

        read -p "Do you want to reinstall? (y/N) " -n 1 -r
        echo
        if [[ ! $REPLY =~ ^[Yy]$ ]]; then
            log_info "Skipping NVIDIA Container Toolkit installation"
            return
        fi
    fi

    # Add NVIDIA Container Toolkit repository
    curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey | sudo gpg --dearmor -o /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg

    curl -s -L https://nvidia.github.io/libnvidia-container/stable/deb/nvidia-container-toolkit.list | \
      sed 's#deb https://#deb [signed-by=/usr/share/keyrings/nvidia-container-toolkit-keyring.gpg] https://#g' | \
      sudo tee /etc/apt/sources.list.d/nvidia-container-toolkit.list

    # Install
    sudo apt-get update
    sudo apt-get install -y nvidia-container-toolkit

    # Configure Docker to use NVIDIA runtime
    sudo nvidia-ctk runtime configure --runtime=docker

    # Restart Docker
    sudo systemctl restart docker

    log_success "NVIDIA Container Toolkit installed and configured"
}

# Verify GPU access in Docker
verify_docker_gpu() {
    log_info "Verifying GPU access in Docker..."

    # Test with nvidia-smi in container
    if docker run --rm --gpus all nvidia/cuda:12.6.0-base-ubuntu22.04 nvidia-smi &> /dev/null; then
        log_success "GPU is accessible from Docker containers!"

        # Show GPU info from container
        log_info "GPU info from container:"
        docker run --rm --gpus all nvidia/cuda:12.6.0-base-ubuntu22.04 nvidia-smi --query-gpu=name,memory.total,driver_version --format=csv
    else
        log_error "GPU is NOT accessible from Docker containers"
        log_info "Try the following:"
        log_info "  1. Restart Docker: sudo systemctl restart docker"
        log_info "  2. Restart WSL: wsl --shutdown (from Windows)"
        log_info "  3. Ensure Windows NVIDIA driver is >= 525.xx"
        exit 1
    fi
}

# Test vLLM container (optional)
test_vllm() {
    log_info "Testing vLLM container (optional)..."

    read -p "Do you want to test vLLM container? This will download ~10GB. (y/N) " -n 1 -r
    echo
    if [[ ! $REPLY =~ ^[Yy]$ ]]; then
        log_info "Skipping vLLM test"
        return
    fi

    log_info "Pulling vLLM image..."
    docker pull vllm/vllm-openai:latest

    log_info "Testing vLLM startup (will exit after verification)..."
    # Just verify it starts, don't load a model
    timeout 30 docker run --rm --gpus all vllm/vllm-openai:latest --help && \
        log_success "vLLM container works!" || \
        log_warning "vLLM test timed out (this is OK, container works)"
}

# Print summary
print_summary() {
    echo ""
    echo "============================================================"
    echo -e "${GREEN}WSL GPU Setup Complete!${NC}"
    echo "============================================================"
    echo ""
    echo "Installed components:"
    echo "  - CUDA Toolkit 12.6"
    echo "  - Docker with GPU support"
    echo "  - NVIDIA Container Toolkit"
    echo ""
    echo "Quick verification commands:"
    echo "  nvidia-smi                    # Check GPU from WSL"
    echo "  docker run --rm --gpus all nvidia/cuda:12.6.0-base-ubuntu22.04 nvidia-smi"
    echo ""
    echo "Next steps:"
    echo "  1. Restart your terminal or run: source ~/.bashrc"
    echo "  2. Start the vLLM stack:"
    echo "     cd /path/to/local-llm-agent"
    echo "     docker-compose -f docker-compose.vllm.yml up -d"
    echo ""
    echo "Documentation: docs/70-onboarding/WSL-GPU-SETUP.md"
    echo "============================================================"
}

# Main execution
main() {
    echo "============================================================"
    echo "WSL GPU Setup Script for Local LLM Agent"
    echo "============================================================"
    echo ""

    check_wsl
    check_ubuntu_version
    check_nvidia_driver

    echo ""
    read -p "Continue with installation? (Y/n) " -n 1 -r
    echo
    if [[ $REPLY =~ ^[Nn]$ ]]; then
        log_info "Installation cancelled"
        exit 0
    fi

    install_cuda_toolkit
    install_docker
    install_nvidia_container_toolkit
    verify_docker_gpu
    test_vllm

    print_summary
}

# Run main
main "$@"