#!/bin/bash
set -e

echo "=== NVIDIA GPU Setup for Docker ==="
echo ""

# Step 1: Detect GPU
echo "=== Step 1: Detecting GPU ==="
GPU_INFO=$(lspci | grep -i nvidia || true)
if [ -z "$GPU_INFO" ]; then
    echo "ERROR: No NVIDIA GPU detected"
    exit 1
fi
echo "$GPU_INFO"
echo ""

# Step 2: Find recommended driver
echo "=== Step 2: Finding recommended driver ==="
RECOMMENDED=$(ubuntu-drivers devices 2>/dev/null | grep "recommended" | awk '{print $3}' || true)
if [ -z "$RECOMMENDED" ]; then
    echo "ERROR: Could not determine recommended driver"
    echo "Available drivers:"
    ubuntu-drivers list
    exit 1
fi
echo "Recommended driver: $RECOMMENDED"
echo ""

# Step 3: Install NVIDIA driver
echo "=== Step 3: Installing $RECOMMENDED ==="
apt-get update
apt-get install -y "$RECOMMENDED"
echo ""

# Step 4: Install NVIDIA Container Toolkit
echo "=== Step 4: Installing NVIDIA Container Toolkit ==="

# Add NVIDIA GPG key
curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey | \
    gpg --dearmor -o /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg --yes

# Add NVIDIA repo
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' | \
    tee /etc/apt/sources.list.d/nvidia-container-toolkit.list

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

# Configure Docker runtime
nvidia-ctk runtime configure --runtime=docker

# Restart Docker
systemctl restart docker
echo ""

# Step 5: Verify installation
echo "=== Step 5: Verifying installation ==="
echo "Installed driver:"
dpkg -l | grep nvidia-driver | head -1
echo ""

echo "============================================"
echo "INSTALLATION COMPLETE"
echo "============================================"
echo ""
echo "IMPORTANT: Reboot required!"
echo ""
echo "After reboot, verify with:"
echo "  nvidia-smi"
echo "  docker run --rm --gpus all ubuntu nvidia-smi"
echo ""
echo "Then run MiniCPM-V 4.5 on GPU:"
echo "  docker run -d --gpus all -p 11434:11434 -v ollama-data:/root/.ollama \\"
echo "    -e MODEL_NAME=\"openbmb/minicpm-v4.5:q8_0\" ht-docker-ai:minicpm45v"
echo ""