#!/bin/bash
set -euo pipefail

# Check bootflag
if ! grep -qw "nvidia=1" /proc/cmdline; then
    echo "Bootflag nvidia=1 not set – skipping."
    exit 0
fi

echo "Bootflag nvidia=1 found – detecting GPU driver type..."

# Determine driver type
driver_type="$(/usr/bin/nvidia-inst --recommended-driver)"

# shopt -s nullglob   # enabling this causes premature exit !

case "$driver_type" in
    nvidia)
        pkg_pattern="/usr/share/packages/nvidia-[0-9]*.pkg.tar.zst"
        ;;
    nvidia-open)
        pkg_pattern="/usr/share/packages/nvidia-open-[0-9]*.pkg.tar.zst"
        ;;
    *)
        echo "${0##*/}: unsupported or unknown driver type '$driver_type' – skipping."
        exit 0
        ;;
esac

pkg_files=($pkg_pattern)

if [ "$pkg_files" = "$pkg_pattern" ]; then
    echo "Error: no packages found for $driver_type"
    exit 1
fi

echo "unloading NVIDIA-Module..."
mod_list=(nvidia_drm nvidia_uvm nvidia_modeset nvidia)
for mod in "${mod_list[@]}"; do
    modprobe -r "$mod" 2>/dev/null || true
done

echo "Installing NVIDIA packages for driver type: $driver_type"
pacman -U --noconfirm "${pkg_files[@]}"

echo "Loading NVIDIA modules..."
modprobe nvidia_drm
modprobe nvidia_uvm
modprobe nvidia_modeset
modprobe nvidia

echo "setting Nvidia GL to be used for the live session"
export __GLX_VENDOR_LIBRARY_NAME=nvidia
export GBM_BACKEND=nvidia-drm

echo "NVIDIA setup done."
