#!/bin/bash
#
# Purpose : Install NVIDIA's driver when the machine has an NVIDIA GPU and no
#           NVIDIA driver yet (i.e. archinstall's proprietary option wasn't
#           picked). Turing (GTX 16xx / RTX 20xx) and newer get the open kernel
#           modules; Maxwell, Pascal and Volta get the 580xx legacy driver from
#           the AUR. Older cards stay on nouveau.
#
#           Run from install.sh only, never from an update: silently moving a
#           working machine off nouveau is too risky to do mid-update.

set -e

source "$HOME/.local/share/coredesktop/utils.sh"

# Read device IDs from sysfs: lspci needs pciutils and wakes suspended GPUs.
# Turing is the first generation with GSP firmware (open modules) and the first
# with device IDs at 0x1e00 or above; Maxwell starts at 0x1340.
GPU=none
for device in /sys/bus/pci/devices/*; do
  [[ $(< "$device/vendor") == 0x10de && $(< "$device/class") == 0x03* ]] || continue
  id=$(( $(< "$device/device") ))
  if (( id >= 0x1e00 )); then
    GPU=open
  elif (( id >= 0x1340 )) && [ "$GPU" = none ]; then
    GPU=legacy
  fi
done

if [ "$GPU" = none ]; then
  echo "No NVIDIA GPU supported by the current drivers found."
  exit 0
fi

if is_installed nvidia-utils || is_installed nvidia-580xx-utils; then
  echo "NVIDIA driver already installed."
  exit 0
fi

# DKMS builds the module for every installed kernel, so each needs its headers
PACKAGES=()
for kernel in linux linux-lts linux-zen linux-hardened; do
  if is_installed "$kernel"; then
    PACKAGES+=("$kernel-headers")
  fi
done

if [ "$GPU" = open ]; then
  PACKAGES+=(nvidia-open-dkms nvidia-utils libva-nvidia-driver)
else
  PACKAGES+=(nvidia-580xx-dkms nvidia-580xx-utils)
fi

echo "Installing NVIDIA driver..."
install_packages "${PACKAGES[@]}"

# Load the driver from the initramfs so the GPU is ready when GDM starts
sudo install -d -m 755 /etc/mkinitcpio.conf.d
echo 'MODULES+=(nvidia nvidia_modeset nvidia_uvm nvidia_drm)' \
  | sudo tee /etc/mkinitcpio.conf.d/nvidia.conf > /dev/null
sudo mkinitcpio -P

# GDM only offers Wayland on NVIDIA when these are enabled, since without them
# the GPU loses its memory contents across suspend
sudo systemctl enable nvidia-suspend.service nvidia-hibernate.service nvidia-resume.service

echo "NVIDIA driver installed. Reboot to start using it."