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

AUTO_REBOOT="${AUTO_REBOOT:-0}"

printf '%s\n' "=== 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

if ! command -v pacman >/dev/null 2>&1; then
  echo "This script is for Arch Linux. pacman was not found."
  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

# Pick NVIDIA package based on the running kernel.
# - linux kernel: nvidia
# - linux-lts kernel: nvidia-lts
# - custom/zen/hardened kernels: nvidia-dkms + matching headers is usually safest
KERNEL_RELEASE="$(uname -r)"
NVIDIA_PACKAGES=(nvidia-utils nvidia-settings opencl-nvidia)

if [[ "$KERNEL_RELEASE" == *-lts ]]; then
  NVIDIA_PACKAGES+=(nvidia-lts)
elif [[ "$KERNEL_RELEASE" == *-arch* ]]; then
  NVIDIA_PACKAGES+=(nvidia)
else
  NVIDIA_PACKAGES+=(nvidia-dkms)
  # Try to install common headers. If this is a custom kernel, install your matching headers manually.
  if pacman -Q linux-headers >/dev/null 2>&1 || pacman -Q linux >/dev/null 2>&1; then
    NVIDIA_PACKAGES+=(linux-headers)
  fi
  if pacman -Q linux-lts >/dev/null 2>&1; then
    NVIDIA_PACKAGES+=(linux-lts-headers)
  fi
  if pacman -Q linux-zen >/dev/null 2>&1; then
    NVIDIA_PACKAGES+=(linux-zen-headers)
  fi
  if pacman -Q linux-hardened >/dev/null 2>&1; then
    NVIDIA_PACKAGES+=(linux-hardened-headers)
  fi
fi

echo
echo "Installing NVIDIA driver packages: ${NVIDIA_PACKAGES[*]}"
pacman -S --needed --noconfirm "${NVIDIA_PACKAGES[@]}"

echo
echo "Enabling NVIDIA persistence daemon if available..."
systemctl enable nvidia-persistenced.service >/dev/null 2>&1 || true

echo
echo "Rebuilding initramfs..."
mkinitcpio -P

echo
echo "Driver install finished. Reboot, then run: nvidia-smi"
echo "If nvidia-smi works, run the second Arch script."

if [[ "$AUTO_REBOOT" == "1" ]]; then
  echo "AUTO_REBOOT=1 set, rebooting now..."
  reboot now
else
  echo "Not rebooting automatically. Run this when ready: sudo reboot"
fi