#!/bin/bash

# ==============================================================================
# Hyprland NVIDIA Setup Script for Arch Linux
# ==============================================================================
# This script automates the installation and configuration of NVIDIA drivers
# for use with Hyprland on Arch Linux, following the official Hyprland wiki.
#
# Author: https://github.com/Kn0ax
# 
# ==============================================================================

# --- Colors for better readability ---
C_RESET='\033[0m'
C_RED='\033[0;31m'
C_YELLOW='\033[0;33m'
C_BOLD='\033[1m'

# --- GPU Detection and Driver Selection ---
# Get GPU info for driver selection
gpu_info=$(lspci | grep -i 'nvidia')

# Turing (16xx, 20xx), Ampere (30xx), Ada (40xx), and newer recommend the open-source kernel modules
if echo "$gpu_info" | grep -q -E "RTX [2-9][0-9]|GTX 16"; then
    NVIDIA_DRIVER_PACKAGE="nvidia-open-dkms"
else
    NVIDIA_DRIVER_PACKAGE="nvidia-dkms"
fi

# Check which kernel is installed and set appropriate headers package
KERNEL_HEADERS="linux-headers"  # Default
if pacman -Q linux-zen &>/dev/null; then
    KERNEL_HEADERS="linux-zen-headers"
elif pacman -Q linux-lts &>/dev/null; then
    KERNEL_HEADERS="linux-lts-headers"
elif pacman -Q linux-hardened &>/dev/null; then
    KERNEL_HEADERS="linux-hardened-headers"
fi

# Install packages
PACKAGES_TO_INSTALL=(
    "${KERNEL_HEADERS}"    
    "${NVIDIA_DRIVER_PACKAGE}"
    "nvidia-utils"
    "lib32-nvidia-utils"
    "egl-wayland"
    "libva-nvidia-driver" # For VA-API hardware acceleration
    "qt5-wayland"
    "qt6-wayland"
)

if ! yay -Syu --needed --noconfirm "${PACKAGES_TO_INSTALL[@]}"; then
    echo -e "${C_RED}${C_BOLD}[ERROR]${C_RESET} Failed to install NVIDIA packages. Please check the output for errors." >&2
    exit 1
fi

# Configure modprobe for early KMS
if ! echo "options nvidia_drm modeset=1" | sudo tee /etc/modprobe.d/nvidia.conf >/dev/null; then
    echo -e "${C_RED}${C_BOLD}[ERROR]${C_RESET} Failed to configure modprobe. Please check your permissions." >&2
    exit 1
fi

# Configure mkinitcpio for early loading
MKINITCPIO_CONF="/etc/mkinitcpio.conf"

if [ ! -f "$MKINITCPIO_CONF" ]; then
    echo -e "${C_RED}${C_BOLD}[ERROR]${C_RESET} $MKINITCPIO_CONF not found. Aborting." >&2
    exit 1
fi

# Define modules
NVIDIA_MODULES="nvidia nvidia_modeset nvidia_uvm nvidia_drm"

# Create backup
sudo cp "$MKINITCPIO_CONF" "${MKINITCPIO_CONF}.backup"

# Remove any old nvidia modules to prevent duplicates
sudo sed -i -E 's/ nvidia_drm//g; s/ nvidia_uvm//g; s/ nvidia_modeset//g; s/ nvidia//g;' "$MKINITCPIO_CONF"
# Add the new modules at the start of the MODULES array
sudo sed -i -E "s/^(MODULES=\\()/\\1${NVIDIA_MODULES} /" "$MKINITCPIO_CONF"
# Clean up potential double spaces
sudo sed -i -E 's/  +/ /g' "$MKINITCPIO_CONF"

if ! sudo mkinitcpio -P; then
    echo -e "${C_RED}${C_BOLD}[ERROR]${C_RESET} Failed to rebuild initramfs with 'mkinitcpio -P'. Please check the output for errors." >&2
    exit 1
fi

# Add NVIDIA environment variables to hyprland.conf
HYPRLAND_CONF="$HOME/.config/hypr/hyprland.conf"
if [ -f "$HYPRLAND_CONF" ]; then
    cat >> "$HYPRLAND_CONF" << 'EOF'

# NVIDIA environment variables
env = ELECTRON_OZONE_PLATFORM_HINT,auto
env = NVD_BACKEND,direct
env = LIBVA_DRIVER_NAME,nvidia
env = __GLX_VENDOR_LIBRARY_NAME,nvidia
EOF
fi

# Final information
echo -e "${C_YELLOW}${C_BOLD}IMPORTANT:${C_RESET} A reboot is required for all changes to take effect."


