#!/usr/bin/env bash
set -euo pipefail
IFS=$'\n\t'
trap 'echo "[ERROR] Line ${LINENO}: command failed" >&2' ERR

echo "[INFO] Detecting NVIDIA GPU..."
if ! lspci | grep -qi 'nvidia'; then
  echo "[FATAL] No NVIDIA GPU detected. This script is for systems with an NVIDIA GPU (e.g. RTX 3060)." >&2
  exit 1
fi

echo "[STEP] Blacklisting nouveau (open-source NVIDIA driver)"
sudo tee /etc/modprobe.d/blacklist-nouveau.conf >/dev/null <<'EOF'
blacklist nouveau
options nouveau modeset=0
EOF

echo "[STEP] Updating initramfs so nouveau stays disabled"
sudo update-initramfs -u

echo "[STEP] Updating APT and upgrading base system"
sudo apt update
sudo apt -y upgrade

echo "[STEP] Installing build/development tools & headers"
sudo apt install -y \
  build-essential \
  dkms \
  linux-headers-$(uname -r) \
  ubuntu-drivers-common

echo "[STEP] Installing recommended NVIDIA driver (with CUDA support)"
# This picks the best proprietary driver for your GPU on Ubuntu 24.04
sudo ubuntu-drivers install --gpgpu

echo
echo "[DONE] NVIDIA proprietary driver install complete."
echo "[NOTE] A reboot is required for the driver to load and nouveau to be disabled."
echo "       Run: sudo reboot"