#!/bin/bash
set -e

echo "--------------------------------------------------------------------"
echo "IMPORTANT: This script will install CUDA 12.8 on Ubuntu 22.04."
echo "Ensure that you have a compatible NVIDIA GPU and that any conflicting CUDA installations are removed."
echo "--------------------------------------------------------------------"
read -p "Are you sure you want to continue? (y/N): " confirmation
if [[ ! "$confirmation" =~ ^[Yy]$ ]]; then
    echo "Aborting."
    exit 1
fi

echo ">>> [1/9] Updating package lists..."
apt update

echo ">>> [2/9] Installing required dependencies..."
apt install -y build-essential dkms wget

echo ">>> [3/9] Downloading CUDA repository pin file..."
wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-ubuntu2204.pin
mv cuda-ubuntu2204.pin /etc/apt/preferences.d/cuda-repository-pin-600

echo ">>> [4/9] Downloading CUDA 12.8 local installer package..."
wget https://developer.download.nvidia.com/compute/cuda/12.8.1/local_installers/cuda-repo-ubuntu2204-12-8-local_12.8.1-570.124.06-1_amd64.deb

echo ">>> [5/9] Installing the CUDA repository package..."
dpkg -i cuda-repo-ubuntu2204-12-8-local_12.8.1-570.124.06-1_amd64.deb

echo ">>> [6/9] Adding the repository GPG key..."
cp /var/cuda-repo-ubuntu2204-12-8-local/cuda-*-keyring.gpg /usr/share/keyrings/

echo ">>> [7/9] Updating package lists after adding CUDA repository..."
apt update

echo ">>> [8/9] Installing CUDA 12.8 with force overwrite..."
apt -o Dpkg::Options::="--force-overwrite" install -y cuda

echo ">>> [9/9] Setting up environment variables for CUDA..."
cat << 'EOF' > /etc/profile.d/cuda.sh
export PATH=/usr/local/cuda-12.8/bin${PATH:+:${PATH}}
export LD_LIBRARY_PATH=/usr/local/cuda-12.8/lib64${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}
EOF
chmod +x /etc/profile.d/cuda.sh

echo ">>> Cleaning up apt cache..."
apt clean

echo ">>> Verifying CUDA installation..."
if command -v nvcc >/dev/null 2>&1; then
    echo "CUDA installed successfully. nvcc version:"
    nvcc --version
else
    echo "nvcc not found. Please check the installation."
fi

echo ""
echo "--------------------------------------------------------------------"
echo "Done! CUDA 12.8 should now be installed on your system."
echo "--------------------------------------------------------------------"

exit 0