#!/bin/bash
# NVIDIA Driver Upgrade Script: 470 -> 535
# Upgrades NVIDIA driver for CUDA 12 support (required by Ollama)

set -e

echo "=== NVIDIA Driver Upgrade: 470 -> 535 ==="
echo ""

# Check if running as root
if [ "$EUID" -ne 0 ]; then
    echo "❌ This script must be run as root (use sudo)"
    exit 1
fi

# Backup current state
echo "[1/9] Backing up current container state..."
docker ps > /tmp/nvidia-upgrade-containers-backup.txt
echo "✅ Container list saved to /tmp/nvidia-upgrade-containers-backup.txt"
echo ""

# Stop GPU containers
echo "[2/9] Stopping GPU-using containers..."
docker stop ollama 2>/dev/null || true
docker stop jellyfin 2>/dev/null || true
echo "✅ GPU containers stopped"
echo ""

# Unhold NVIDIA Container Toolkit packages
echo "[3/9] Unholding NVIDIA Container Toolkit packages..."
echo "libnvidia-container-tools install" | dpkg --set-selections
echo "libnvidia-container1 install" | dpkg --set-selections
echo "nvidia-container-toolkit install" | dpkg --set-selections
echo "nvidia-container-toolkit-base install" | dpkg --set-selections
echo "✅ Packages unheld"
echo ""

# Remove ALL driver 470 packages
echo "[4/9] Removing ALL NVIDIA driver 470 packages..."
apt-get purge -y 'libnvidia-*-470' 'nvidia-*-470' '*nvidia*470*' 2>&1 | grep -E "Removing|Purging|^$" || true
apt-get purge -y nvidia-settings 2>&1 | grep -E "Removing|Purging|^$" || true
apt-get autoremove -y > /dev/null 2>&1
echo "✅ All driver 470 packages removed"
echo ""

# Clean up package database
echo "[5/9] Cleaning package database..."
apt-get clean
apt-get autoclean
dpkg --configure -a
apt-get install -f -y
echo "✅ Package database cleaned"
echo ""

# Update package lists
echo "[6/9] Updating package lists..."
apt-get update > /dev/null 2>&1
echo "✅ Package lists updated"
echo ""

# Install driver 535 with all dependencies
echo "[7/9] Installing NVIDIA driver 535..."
apt-get install -y --no-install-recommends \
    libnvidia-common-535 \
    libnvidia-gl-535 \
    libnvidia-compute-535 \
    libnvidia-decode-535 \
    libnvidia-encode-535 \
    nvidia-utils-535 \
    xserver-xorg-video-nvidia-535 \
    libnvidia-cfg1-535 \
    libnvidia-fbc1-535 \
    nvidia-kernel-common-535 \
    nvidia-dkms-535 \
    nvidia-driver-535 \
    nvidia-settings

echo "✅ Driver 535 installed"
echo ""

# Update NVIDIA Container Toolkit
echo "[8/9] Updating NVIDIA Container Toolkit..."
apt-get install --reinstall -y nvidia-container-toolkit > /dev/null 2>&1
nvidia-ctk runtime configure --runtime=docker > /dev/null 2>&1
echo "✅ Container toolkit updated"
echo ""

# Restart Docker
echo "[9/9] Restarting Docker daemon..."
systemctl restart docker
sleep 3
echo "✅ Docker restarted"
echo ""

echo "=== Upgrade Complete ==="
echo ""
echo "⚠️  REBOOT REQUIRED NOW"
echo ""
echo "Run: sudo reboot"
echo ""
echo "After reboot:"
echo "  1. Verify:       nvidia-smi"
echo "  2. Restart:      cd ~/Projects/portainer-core/stacks && docker compose -f ollama.yml up -d"
echo "  3. Test GPU:     docker exec ollama ollama ps"
echo ""