#!/usr/bin/env bash

set -euo pipefail
IFS=$'\n\t'

LOGFILE=${LOGFILE:-/var/log/nvidia-driver-install.log}
exec > >(tee -a "$LOGFILE") 2>&1

trap 'echo "[ERROR] Line ${LINENO}: command failed"; exit 1' ERR

require_root() {
  if [[ ${EUID:-$(id -u)} -ne 0 ]]; then
    echo "[INFO] Re-running with sudo..."
    exec sudo -E bash "$0" "$@"
  fi
}
require_root "$@"

# Basic platform sanity
if ! command -v apt >/dev/null 2>&1; then
  echo "[FATAL] This script targets APT-based distros (Ubuntu/Debian)."
  exit 2
fi

# Secure Boot check (unsigned DKMS modules will not load)
if command -v mokutil >/dev/null 2>&1; then
  if mokutil --sb-state 2>/dev/null | grep -qi 'enabled'; then
    echo "[WARN] Secure Boot appears ENABLED. NVIDIA kernel modules may fail to load."
    echo "       Consider disabling Secure Boot or enrolling MOK keys before proceeding."
  fi
fi

echo "[STEP] Updating APT metadata"
apt-get update -y

echo "[STEP] Fixing broken packages (if any)"
apt-get -y --fix-broken install || true

echo "[STEP] Purging conflicting server-flavor NVIDIA packages (if present)"
dpkg -l | awk '/nvidia-.*server/ {print $2}' | xargs -r apt-get purge -y

echo "[STEP] Installing prerequisites"
DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends       dkms build-essential linux-headers-$(uname -r) pciutils curl ca-certificates mokutil || true

# Prefer ubuntu-drivers if available (Ubuntu)
if command -v ubuntu-drivers >/dev/null 2>&1; then
  echo "[STEP] Using ubuntu-drivers to select the recommended driver"
  ubuntu-drivers autoinstall
else
  # Fallback to meta-driver if ubuntu-drivers is not available
  echo "[STEP] Installing NVIDIA driver meta-package"
  # Try a reasonable meta; the specific version changes often, so use the generic meta when available
  apt-get install -y nvidia-driver || apt-get install -y nvidia-driver-535 || true
fi

echo "[STEP] Cleaning up"
apt-get -y autoremove
apt-get -y autoclean

echo "[INFO] Verifying installation"
if command -v nvidia-smi >/dev/null 2>&1; then
  nvidia-smi || true
else
  echo "[WARN] nvidia-smi not found yet. A reboot may be required."
fi

echo "[DONE] NVIDIA driver install complete. Reboot is recommended."