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

if [[ $EUID -ne 0 ]]; then
  echo "❌ Run as root: sudo $0"
  exit 1
fi

echo "=== NVIDIA driver install (Ubuntu 24.04) ==="

apt-get update
apt-get install -y ubuntu-drivers-common

echo
echo "📋 ubuntu-drivers devices:"
ubuntu-drivers devices || true

echo
read -rp "👉 Install recommended NVIDIA driver via 'ubuntu-drivers autoinstall'? [y/N]: " CONFIRM
if [[ "$CONFIRM" =~ ^[Yy]$ ]]; then
  ubuntu-drivers autoinstall
  echo "✅ Driver install finished."
else
  echo "⏭️ Skipped."
  exit 0
fi

echo
read -rp "🔁 Reboot now? [y/N]: " REBOOT
if [[ "$REBOOT" =~ ^[Yy]$ ]]; then
  reboot
else
  echo "⚠️ Please reboot manually before running stage 2."
fi
