#!/bin/bash

# Ensure the script is run as root
if [[ $EUID -ne 0 ]]; then
    echo "This script must be run as root. Try: sudo $0"
    exit 1
fi

# Update and upgrade the system
echo "Updating system..."
apt update && apt upgrade -y

# Install necessary dependencies
echo "Installing required packages..."
apt install -y build-essential dkms linux-headers-$(uname -r)

# Detect and install NVIDIA driver
echo "Detecting GPU model..."
apt install -y pciutils
GPU_MODEL=$(lspci | grep -i NVIDIA)

if [[ -z "$GPU_MODEL" ]]; then
    echo "No NVIDIA GPU detected. Exiting..."
    exit 1
fi

echo "Installing NVIDIA driver..."
apt install -y nvidia-driver

# Reboot prompt
echo "Installation complete! A system reboot is recommended."
read -p "Do you want to reboot now? (y/n) " REBOOT

if [[ "$REBOOT" == "y" ]]; then
    reboot
fi

exit 0