#!/usr/bin/env bash
# Ubuntu 24.04.x (Noble) + RTX 5070 Ti
# Target state (simple/reliable):
# - NVIDIA driver: Ubuntu packages (nvidia-driver-580-open + libs)
# - CUDA Toolkit: NVIDIA repo (cuda-toolkit-12-8) under /usr/local/cuda-12.8
# - APT pinning: prevent NVIDIA CUDA repo from “winning” for nvidia-* / libnvidia-* packages
#
# Run:
#   chmod +x ~/setup-rtx5070ti-ubuntu2404.sh
#   sudo ~/setup-rtx5070ti-ubuntu2404.sh

set -euo pipefail

LOG="/tmp/rtx5070ti-ubuntu2404-setup.$(date +%Y%m%d_%H%M%S).log"
exec > >(tee -a "$LOG") 2>&1

fatal() { echo "ERROR: $*" >&2; exit 1; }
info()  { echo "INFO:  $*"; }
warn()  { echo "WARN:  $*"; }

require_cmd() {
  command -v "$1" >/dev/null 2>&1 || fatal "Missing required command: $1"
}

require_root() {
  if [[ "${EUID:-$(id -u)}" -ne 0 ]]; then
    fatal "Run as root (e.g. sudo $0)"
  fi
}

check_os() {
  [[ -r /etc/os-release ]] || fatal "/etc/os-release not found"
  # shellcheck disable=SC1091
  . /etc/os-release
  [[ "${ID:-}" == "ubuntu" ]] || fatal "Ubuntu required (detected ID=${ID:-unknown})"
  [[ "${VERSION_ID:-}" == "24.04" ]] || fatal "Ubuntu 24.04.x required (detected VERSION_ID=${VERSION_ID:-unknown})"
  info "Detected Ubuntu ${VERSION_ID} (${VERSION_CODENAME:-unknown})."
}

check_secure_boot() {
  if command -v mokutil >/dev/null 2>&1; then
    local sb
    sb="$(mokutil --sb-state 2>/dev/null || true)"
    if echo "$sb" | grep -qi "enabled"; then
      warn "Secure Boot appears ENABLED. If NVIDIA modules fail to load, Secure Boot is a common cause."
      warn "For a deterministic first-boot experience, Secure Boot OFF is recommended."
    fi
  fi
}

write_apt_pinning() {
  local pref="/etc/apt/preferences.d/99-nvidia-from-ubuntu-cuda-from-nvidia.pref"
  info "Configuring APT pinning at: $pref"
  cat > "$pref" <<'PREF'
# Keep NVIDIA driver & libs from Ubuntu, not from the NVIDIA CUDA repo.
Package: nvidia-*
Pin: origin "developer.download.nvidia.com"
Pin-Priority: 1

Package: libnvidia-*
Pin: origin "developer.download.nvidia.com"
Pin-Priority: 1

Package: linux-modules-nvidia-*
Pin: origin "developer.download.nvidia.com"
Pin-Priority: 1

Package: xserver-xorg-video-nvidia-*
Pin: origin "developer.download.nvidia.com"
Pin-Priority: 1

# Keep CUDA Toolkit from the NVIDIA repo.
Package: cuda-*
Pin: origin "developer.download.nvidia.com"
Pin-Priority: 600
PREF
}

apt_update_upgrade() {
  export DEBIAN_FRONTEND=noninteractive
  info "APT update + base upgrade"
  apt-get update -y
  apt-get upgrade -y
}

install_driver_ubuntu() {
  export DEBIAN_FRONTEND=noninteractive
  local krel
  krel="$(uname -r)"
  info "Installing NVIDIA driver from Ubuntu (580-open) for kernel: $krel"

  apt-get install -y \
    "linux-modules-nvidia-580-open-$krel" \
    nvidia-driver-580-open \
    nvidia-utils-580 \
    libnvidia-gl-580

  update-initramfs -u
}

remove_ubuntu_cuda_toolkit_if_present() {
  export DEBIAN_FRONTEND=noninteractive
  if dpkg -l | awk '{print $2}' | grep -qx "nvidia-cuda-toolkit"; then
    info "Removing Ubuntu 'nvidia-cuda-toolkit' to avoid mixing toolkits"
    apt-get purge -y nvidia-cuda-toolkit
    apt-get autoremove -y
  else
    info "Ubuntu 'nvidia-cuda-toolkit' not installed (good)."
  fi
}

add_nvidia_cuda_repo_if_missing() {
  export DEBIAN_FRONTEND=noninteractive
  local keypkg="/tmp/cuda-keyring_1.1-1_all.deb"

  # If already present, do nothing.
  if grep -Rqs "developer.download.nvidia.com/compute/cuda/repos/ubuntu2404" /etc/apt/sources.list /etc/apt/sources.list.d/* 2>/dev/null; then
    info "NVIDIA CUDA repo already configured."
    return
  fi

  info "Adding NVIDIA CUDA repo (cuda-keyring)"
  apt-get update -y
  apt-get install -y wget gnupg ca-certificates

  wget -q -O "$keypkg" "https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2404/x86_64/cuda-keyring_1.1-1_all.deb"
  dpkg -i "$keypkg"
}

install_cuda_12_8() {
  export DEBIAN_FRONTEND=noninteractive
  info "Installing CUDA Toolkit 12.8 from NVIDIA repo"
  apt-get update -y
  apt-get install -y cuda-toolkit-12-8
}

configure_cuda_env_systemwide() {
  local f="/etc/profile.d/cuda-12-8.sh"
  info "Writing system-wide CUDA environment file: $f"
  cat > "$f" <<'ENVF'
export PATH=/usr/local/cuda-12.8/bin:$PATH
export LD_LIBRARY_PATH=/usr/local/cuda-12.8/lib64:$LD_LIBRARY_PATH
ENVF

  # Remove duplicates from the invoking user's ~/.bashrc if present.
  local u="${SUDO_USER:-}"
  if [[ -n "$u" && -d "/home/$u" && -f "/home/$u/.bashrc" ]]; then
    info "Cleaning CUDA 12.8 exports from /home/$u/.bashrc (single source of truth: /etc/profile.d)"
    sed -i '/\/usr\/local\/cuda-12\.8\/bin/d;/\/usr\/local\/cuda-12\.8\/lib64/d' "/home/$u/.bashrc" || true
  fi
}

validate_now() {
  info "Validation (current boot):"

  if command -v nvidia-smi >/dev/null 2>&1; then
    nvidia-smi || warn "nvidia-smi failed (often resolved after reboot if modules were just installed)."
  else
    warn "nvidia-smi not found (unexpected if nvidia-utils-580 installed)."
  fi

  if [[ -x /usr/local/cuda-12.8/bin/nvcc ]]; then
    /usr/local/cuda-12.8/bin/nvcc --version || warn "nvcc failed to run."
  else
    warn "nvcc not found at /usr/local/cuda-12.8/bin/nvcc (CUDA install may have failed)."
  fi

  info "APT policy sanity check (candidates should prefer Ubuntu for driver/libs):"
  apt-cache policy libnvidia-gl-580 libnvidia-egl-gbm1 nvidia-driver-580-open nvidia-utils-580 | sed -n '1,220p' || true
}

prompt_reboot() {
  echo
  info "Recommended: reboot now to load NVIDIA kernel modules cleanly."
  read -r -p "Reboot now? [y/N] " ans
  if [[ "${ans,,}" == "y" || "${ans,,}" == "yes" ]]; then
    info "Rebooting..."
    reboot
  else
    warn "You chose not to reboot now. If anything looks off, reboot and re-check with: nvidia-smi ; nvcc --version"
    info "Log saved to: $LOG"
  fi
}

main() {
  require_root
  require_cmd apt-get
  require_cmd dpkg
  require_cmd uname
  require_cmd update-initramfs

  check_os
  check_secure_boot

  # Pinning first so the NVIDIA CUDA repo cannot override the driver/libs later.
  write_apt_pinning

  apt_update_upgrade
  install_driver_ubuntu

  # CUDA: keep NVIDIA 12.8 only (remove Ubuntu toolkit if present).
  remove_ubuntu_cuda_toolkit_if_present
  add_nvidia_cuda_repo_if_missing
  install_cuda_12_8
  configure_cuda_env_systemwide

  # Repair any partial states and finish upgrades.
  info "Final APT fix-broken (if needed) + upgrade"
  apt-get -f install -y || true
  apt-get upgrade -y

  validate_now
  prompt_reboot
}

main "$@"