#!/usr/bin/env bash
set -euo pipefail

echo "=== Step 1: NVIDIA driver install for RTX 3060 / RTX 3090 on Arch Linux ==="

if [[ $EUID -ne 0 ]]; then
  echo "Please run this script as root, e.g.: sudo bash $0"
  exit 1
fi

echo "Updating Arch package database and system..."
pacman -Syu --noconfirm

echo
echo "Installing base dependencies..."
pacman -S --needed --noconfirm \
  curl \
  ca-certificates \
  gnupg \
  lsb-release \
  git \
  base-devel \
  linux-firmware

# Install kernel headers for whichever common Arch kernels are installed.
# DKMS needs matching headers to build the NVIDIA kernel module.
HEADER_PKGS=()
pacman -Q linux >/dev/null 2>&1 && HEADER_PKGS+=(linux-headers)
pacman -Q linux-lts >/dev/null 2>&1 && HEADER_PKGS+=(linux-lts-headers)
pacman -Q linux-zen >/dev/null 2>&1 && HEADER_PKGS+=(linux-zen-headers)
pacman -Q linux-hardened >/dev/null 2>&1 && HEADER_PKGS+=(linux-hardened-headers)

if (( ${#HEADER_PKGS[@]} > 0 )); then
  echo
  echo "Installing kernel headers: ${HEADER_PKGS[*]}"
  pacman -S --needed --noconfirm "${HEADER_PKGS[@]}"
else
  echo "WARNING: No standard Arch kernel package detected."
  echo "If you use a custom kernel, install its matching headers before rebooting."
fi

echo
echo "Installing NVIDIA driver packages for current Arch repos..."
echo "Using nvidia-open-dkms instead of old/removed nvidia package."
pacman -S --needed --noconfirm \
  nvidia-open-dkms \
  nvidia-utils \
  nvidia-settings \
  opencl-nvidia \
  lib32-nvidia-utils \
  lib32-opencl-nvidia \
  egl-wayland \
  vulkan-icd-loader \
  lib32-vulkan-icd-loader

# Make sure nouveau does not grab the card first.
echo
echo "Blacklisting nouveau..."
cat >/etc/modprobe.d/blacklist-nouveau.conf <<'EONOUVEAU'
blacklist nouveau
options nouveau modeset=0
EONOUVEAU

# Enable NVIDIA DRM KMS for Wayland/KDE and smoother boot.
echo
echo "Enabling NVIDIA DRM modeset..."
cat >/etc/modprobe.d/nvidia-drm.conf <<'EONVIDIA'
options nvidia_drm modeset=1 fbdev=1
EONVIDIA

# Rebuild initramfs if mkinitcpio exists.
if command -v mkinitcpio >/dev/null 2>&1; then
  echo
  echo "Rebuilding initramfs..."
  mkinitcpio -P
fi

echo
echo "Driver install finished. Reboot, then run:"
echo "  nvidia-smi"
echo
echo "If nvidia-smi shows your RTX 3060/3090, continue with the second script."