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

echo "=== Smart NVIDIA driver install for RTX 3060 / RTX 3090 on Arch Linux ==="

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

REALUSER="${SUDO_USER:-$USER}"

run() {
  echo "+ $*"
  "$@"
}

pkg_installed() {
  pacman -Qq "$1" >/dev/null 2>&1
}

pkg_available() {
  pacman -Si "$1" >/dev/null 2>&1
}

any_pkg_installed() {
  local p
  for p in "$@"; do
    if pkg_installed "$p"; then return 0; fi
  done
  return 1
}

kernel_pkgs=()
header_pkgs=()

# Detect installed Arch kernels and match headers.
if pkg_installed linux; then
  kernel_pkgs+=(linux)
  header_pkgs+=(linux-headers)
fi
if pkg_installed linux-lts; then
  kernel_pkgs+=(linux-lts)
  header_pkgs+=(linux-lts-headers)
fi
if pkg_installed linux-zen; then
  kernel_pkgs+=(linux-zen)
  header_pkgs+=(linux-zen-headers)
fi
if pkg_installed linux-hardened; then
  kernel_pkgs+=(linux-hardened)
  header_pkgs+=(linux-hardened-headers)
fi

if [[ ${#kernel_pkgs[@]} -eq 0 ]]; then
  echo "WARNING: Could not detect a standard Arch kernel package."
  echo "This script will use DKMS because it is safest for custom kernels."
fi

echo
echo "Detected kernel packages: ${kernel_pkgs[*]:-(none detected)}"
echo "Required header packages: ${header_pkgs[*]:-(none detected)}"
echo

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

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

if [[ ${#header_pkgs[@]} -gt 0 ]]; then
  echo
  echo "Installing matching kernel headers..."
  run pacman -S --needed --noconfirm "${header_pkgs[@]}"
fi

echo
echo "Checking NVIDIA GPU model..."
GPU_LINES="$(lspci -nn | grep -Ei 'vga|3d|display' | grep -Ei 'nvidia|geforce|rtx|gtx' || true)"
if [[ -z "$GPU_LINES" ]]; then
  echo "WARNING: No NVIDIA GPU found by lspci. Continuing anyway."
else
  echo "$GPU_LINES"
fi

# RTX 3060 and RTX 3090 are Ampere GPUs. Current Arch packages use open kernel modules
# for Turing and newer. Pascal/Maxwell need legacy 580xx from AUR, but that is NOT your GPU.
if echo "$GPU_LINES" | grep -Eiq 'RTX 3060|RTX 3090|GA10|Ampere|NVIDIA.*30[0-9][0-9]'; then
  echo "Detected RTX 30-series/Ampere-class GPU: correct target is current nvidia-open stack."
else
  echo "NOTE: I did not positively identify an RTX 3060/3090 from lspci text."
  echo "If this machine has a 3060/3090 installed, the nvidia-open stack is still correct."
fi

# Choose driver flavor intelligently.
# - If nvidia-open is already installed, keep it unless non-standard kernels are installed without linux.
# - If multiple/alternate kernels are installed, DKMS is more flexible.
# - Do not install nvidia-open and nvidia-open-dkms together; they conflict.
DRIVER_PKG=""

if pkg_installed nvidia-open; then
  DRIVER_PKG="nvidia-open"
  echo
  echo "Existing driver detected: nvidia-open. Keeping it."
elif pkg_installed nvidia-open-dkms; then
  DRIVER_PKG="nvidia-open-dkms"
  echo
  echo "Existing driver detected: nvidia-open-dkms. Keeping it."
else
  if [[ ${#kernel_pkgs[@]} -eq 1 && " ${kernel_pkgs[*]} " == *" linux "* ]]; then
    DRIVER_PKG="nvidia-open"
  else
    DRIVER_PKG="nvidia-open-dkms"
  fi
  echo
  echo "No existing NVIDIA open driver package detected. Selected: ${DRIVER_PKG}"
fi

# If the selected driver conflicts with an installed variant, remove only the conflicting variant.
if [[ "$DRIVER_PKG" == "nvidia-open" ]] && pkg_installed nvidia-open-dkms; then
  echo "Removing conflicting package nvidia-open-dkms before installing nvidia-open..."
  run pacman -Rns --noconfirm nvidia-open-dkms
fi

if [[ "$DRIVER_PKG" == "nvidia-open-dkms" ]] && pkg_installed nvidia-open; then
  echo "Removing conflicting package nvidia-open before installing nvidia-open-dkms..."
  run pacman -Rns --noconfirm nvidia-open
fi

# Remove old proprietary/legacy packages if present and conflicting.
old_pkgs=()
for p in nvidia nvidia-dkms nvidia-lts nvidia-580xx-dkms nvidia-580xx-utils nvidia-580xx-settings; do
  if pkg_installed "$p"; then old_pkgs+=("$p"); fi
done
if [[ ${#old_pkgs[@]} -gt 0 ]]; then
  echo "Removing old/conflicting NVIDIA packages: ${old_pkgs[*]}"
  run pacman -Rns --noconfirm "${old_pkgs[@]}" || true
fi

echo
echo "Installing NVIDIA userspace + selected kernel module package..."
packages=("$DRIVER_PKG" nvidia-utils nvidia-settings opencl-nvidia egl-wayland vulkan-icd-loader)

if pacman -Si lib32-nvidia-utils >/dev/null 2>&1; then
  packages+=(lib32-nvidia-utils lib32-vulkan-icd-loader)
fi

if [[ "$DRIVER_PKG" == "nvidia-open-dkms" ]]; then
  packages+=(dkms)
fi

run pacman -S --needed --noconfirm "${packages[@]}"

echo
echo "Blacklisting nouveau just in case..."
cat >/etc/modprobe.d/blacklist-nouveau.conf <<'EON'
blacklist nouveau
options nouveau modeset=0
EON

# Enable early KMS for NVIDIA.
echo
echo "Configuring NVIDIA DRM modeset..."
cat >/etc/modprobe.d/nvidia.conf <<'EON'
options nvidia_drm modeset=1 fbdev=1
EON

# Rebuild initramfs.
echo
echo "Rebuilding initramfs..."
run mkinitcpio -P

# Enable persistence daemon if available.
if systemctl list-unit-files | grep -q '^nvidia-persistenced.service'; then
  echo
  echo "Enabling NVIDIA persistence daemon..."
  run systemctl enable --now nvidia-persistenced.service || true
fi

echo
echo "Installed NVIDIA packages:"
pacman -Q | grep -E '^(nvidia|lib32-nvidia|opencl-nvidia|egl-wayland|vulkan-icd-loader|dkms) ' || true

echo
echo "Kernel in use: $(uname -r)"
echo
echo "Current NVIDIA module status before reboot:"
if command -v nvidia-smi >/dev/null 2>&1; then
  nvidia-smi || true
else
  echo "nvidia-smi not found yet."
fi

echo
echo "==================== DONE ===================="
echo "Driver package selected: ${DRIVER_PKG}"
echo "For RTX 3060 / RTX 3090 on current Arch, nvidia-open or nvidia-open-dkms is correct."
echo "Reboot now, then run: nvidia-smi"
echo "=============================================="