#!/bin/bash
# ==========================================================
# ProxMenux - NVIDIA GPU Driver Installer
# ==========================================================
# Author      : MacRimi
# Copyright   : (c) 2024 MacRimi
# License     : GPL-3.0
# Version     : 1.2
# Last Updated: 26/03/2026
# ==========================================================
# Description:
# Installs and manages the NVIDIA proprietary driver on a
# Proxmox VE host. Detects hardware, picks a kernel-compatible
# driver version and handles the full lifecycle
# (install / update / remove).
#
# Features:
#  - GPU detection + VFIO passthrough safety check
#  - Kernel-aware driver version filter (5.15 → 6.17+)
#  - Nouveau blacklist + module unload
#  - DKMS-backed install (survives kernel upgrades)
#  - udev rules + nvidia-persistenced service
#  - Optional keylase/nvidia-patch (NVENC session limit)
#  - LXC container driver propagation (Alpine/Arch/Debian)
#  - Complete uninstall path
# ==========================================================

SCRIPT_TITLE="NVIDIA GPU Driver Installer for Proxmox VE"

LOCAL_SCRIPTS="/usr/local/share/proxmenux/scripts"
BASE_DIR="/usr/local/share/proxmenux"
UTILS_FILE="$BASE_DIR/utils.sh"
COMPONENTS_STATUS_FILE="$BASE_DIR/components_status.json"
LOG_FILE="/tmp/nvidia_install.log"
screen_capture="/tmp/proxmenux_nvidia_screen_capture_$$.txt"

NVIDIA_BASE_URL="https://download.nvidia.com/XFree86/Linux-x86_64"
NVIDIA_WORKDIR="/opt/nvidia"

# LXC post-install update constants (used only when NVIDIA LXC passthrough
# containers are detected and the user confirms updating them after the host
# install/reinstall finishes).
NVIDIA_INSTALL_MIN_MB=2048
CT_ORIG_MEM=""

export BASE_DIR
export COMPONENTS_STATUS_FILE

if [[ -f "$UTILS_FILE" ]]; then
  source "$UTILS_FILE"
fi
if [[ -f "$LOCAL_SCRIPTS/global/pci_passthrough_helpers.sh" ]]; then
  source "$LOCAL_SCRIPTS/global/pci_passthrough_helpers.sh"
elif [[ -f "$(cd "$(dirname "${BASH_SOURCE[0]}")"/.. && pwd)/global/pci_passthrough_helpers.sh" ]]; then
  source "$(cd "$(dirname "${BASH_SOURCE[0]}")"/.. && pwd)/global/pci_passthrough_helpers.sh"
fi

if [[ ! -f "$COMPONENTS_STATUS_FILE" ]]; then
  echo "{}" > "$COMPONENTS_STATUS_FILE"
fi

load_language
initialize_cache

# ==========================================================
# GPU detection and current status
# ==========================================================
# Populated by detect_nvidia_gpus. Holds every video-controller PCI
# Device ID (lowercase, 4-hex) so the version filter can drop branches
# whose supportedchips.html doesn't list every card on this host.
NVIDIA_HOST_GPU_IDS=()

detect_nvidia_gpus() {
  # Video controllers only — the paired HDA audio functions (10de:xxxx
  # under class 0403) are not what the display driver ships support for.
  local lspci_output
  lspci_output=$(lspci -nn | grep -i "NVIDIA" \
    | grep -Ei "VGA compatible controller|3D controller|Display controller" || true)

  if [[ -z "$lspci_output" ]]; then
    NVIDIA_GPU_PRESENT=false
    DETECTED_GPUS_TEXT="$(translate 'No NVIDIA GPU detected on this system.')"
    NVIDIA_HOST_GPU_IDS=()
  else
    NVIDIA_GPU_PRESENT=true
    DETECTED_GPUS_TEXT=""
    NVIDIA_HOST_GPU_IDS=()
    local i=1
    while IFS= read -r line; do
      DETECTED_GPUS_TEXT+="  ${i}. ${line}\n"
      # Extract [10de:XXXX] — Vendor:Device pair. We keep only the
      # Device half (4-hex) lowercased, which is what NVIDIA lists in
      # each version's README/supportedchips.html.
      local dev_id
      dev_id=$(echo "$line" | grep -oiE '\[10de:[0-9a-f]{4}\]' | head -1 \
        | sed -E 's/^\[10de:([0-9a-f]{4})\]$/\1/i' | tr 'A-F' 'a-f')
      if [[ -n "$dev_id" ]]; then
        NVIDIA_HOST_GPU_IDS+=("$dev_id")
      fi
      ((i++))
    done <<< "$lspci_output"
  fi
}

check_gpu_not_in_vm_passthrough() {
  local dev vendor driver vfio_list=""
  for dev in /sys/bus/pci/devices/*; do
    vendor=$(cat "$dev/vendor" 2>/dev/null)
    [[ "$vendor" != "0x10de" ]] && continue
    if [[ -L "$dev/driver" ]]; then
      driver=$(basename "$(readlink "$dev/driver")")
      if [[ "$driver" == "vfio-pci" ]]; then
        vfio_list+="  • $(basename "$dev")\n"
      fi
    fi
  done

  [[ -z "$vfio_list" ]] && return 0

  local msg
  msg="\n$(translate "One or more NVIDIA GPUs are currently configured for VM passthrough (vfio-pci):")\n\n"
  msg+="${vfio_list}\n"
  msg+="$(translate "Installing host drivers while the GPU is assigned to a VM could break passthrough and destabilize the system.")\n\n"
  msg+="$(translate "To install host drivers, first remove the GPU from VM passthrough configuration and reboot.")"

  dialog --backtitle "ProxMenux" \
    --title "$(translate "GPU in VM Passthrough Mode")" \
    --msgbox "$msg" 16 78
  exit 0
}

check_stale_vfio_config_for_nvidia() {
  local vfio_conf="/etc/modprobe.d/vfio.conf"
  [[ ! -f "$vfio_conf" ]] && return 0

  local ids_line ids_part
  ids_line=$(grep "^options vfio-pci ids=" "$vfio_conf" 2>/dev/null | head -1)
  [[ -z "$ids_line" ]] && return 0
  ids_part=$(echo "$ids_line" | grep -oE 'ids=[^[:space:]]+' | sed 's/ids=//')
  [[ -z "$ids_part" ]] && return 0

  local dev vendor did vid_did
  local -a legacy_ids=()
  local legacy_list=""

  for dev in /sys/bus/pci/devices/*; do
    vendor=$(cat "$dev/vendor" 2>/dev/null)
    [[ "$vendor" != "0x10de" ]] && continue
    did=$(cat "$dev/device" 2>/dev/null)
    [[ -z "$did" ]] && continue
    vid_did="10de:${did#0x}"
    if echo ",${ids_part}," | grep -q ",${vid_did},"; then
      legacy_ids+=("$vid_did")
      legacy_list+="  • $(basename "$dev")  [${vid_did}]\n"
    fi
  done

  [[ ${#legacy_ids[@]} -eq 0 ]] && return 0

  local msg
  msg="\n$(translate 'A previous VFIO passthrough configuration was detected for the following NVIDIA GPU(s):')\n\n"
  msg+="${legacy_list}\n"
  msg+="$(translate 'The active kernel driver is not vfio-pci, but the entry in') /etc/modprobe.d/vfio.conf $(translate 'will rebind the GPU to vfio-pci on the next reboot, breaking the driver that is about to be installed.')\n\n"
  msg+="\Z1\Zb$(translate 'Do you want to remove the stale entry from vfio.conf and continue?')\Zn"

  dialog --colors --backtitle "ProxMenux" \
    --title "$(translate 'Stale VFIO Config Detected')" \
    --yesno "$msg" 18 78 || exit 0

  if declare -F _clean_vfio_conf_ids >/dev/null 2>&1 \
     && _clean_vfio_conf_ids "${legacy_ids[@]}"; then
    msg_info "$(translate 'Rebuilding initramfs after vfio.conf cleanup...')"
    update-initramfs -u >/dev/null 2>&1 || true
    msg_ok "$(translate 'Stale VFIO entries removed and initramfs rebuilt.')" | tee -a "$screen_capture"
  else
    msg_ok "$(translate 'No changes were needed in vfio.conf.')" | tee -a "$screen_capture"
  fi
}

detect_driver_status() {
  CURRENT_DRIVER_INSTALLED=false
  CURRENT_DRIVER_VERSION=""
  
  # First check if nvidia kernel module is actually loaded
  if grep -q "^nvidia " /proc/modules 2>/dev/null; then

    modprobe nvidia-uvm 2>/dev/null || true
    sleep 1
    

    if command -v nvidia-smi >/dev/null 2>&1; then
      CURRENT_DRIVER_VERSION=$(nvidia-smi --query-gpu=driver_version --format=csv,noheader 2>/dev/null | head -n1)
      
      if [[ -n "$CURRENT_DRIVER_VERSION" ]]; then
        CURRENT_DRIVER_INSTALLED=true
        # Register the installed driver version in components_status.json
        update_component_status "nvidia_driver" "installed" "$CURRENT_DRIVER_VERSION" "gpu" '{"patched":false}'
      fi
    fi
  fi

  if $CURRENT_DRIVER_INSTALLED; then
    CURRENT_STATUS_TEXT="$(printf '%s %s' "$(translate 'NVIDIA driver installed:')" "$CURRENT_DRIVER_VERSION")"
  else
    CURRENT_STATUS_TEXT="$(translate 'No NVIDIA driver installed.')"
  fi

  if $CURRENT_DRIVER_INSTALLED; then
    CURRENT_STATUS_COLORED="${CURRENT_STATUS_TEXT}"
  else
    CURRENT_STATUS_COLORED="${CURRENT_STATUS_TEXT}"
  fi
}

# ==========================================================
# LXC NVIDIA passthrough — discovery & userspace-libs update
# Invoked after the host install/reinstall completes. Aligned with the install
# path used in add_gpu_lxc.sh (distro-aware, memory/disk checks, --no-dkms,
# --no-install-compat32-libs, visible progress via tee).
# ==========================================================
find_nvidia_containers() {
  NVIDIA_CONTAINERS=()
  for conf in /etc/pve/lxc/*.conf; do
    [[ -f "$conf" ]] || continue
    if grep -qiE "dev[0-9]+:.*nvidia" "$conf"; then
      NVIDIA_CONTAINERS+=("$(basename "$conf" .conf)")
    fi
  done
}

get_lxc_nvidia_version() {
  local ctid="$1"
  local version=""

  # Prefer nvidia-smi when the container is running (works with .run-installed drivers)
  if pct status "$ctid" 2>/dev/null | grep -q "running"; then
    version=$(pct exec "$ctid" -- nvidia-smi \
      --query-gpu=driver_version --format=csv,noheader 2>/dev/null \
      | head -1 | tr -d '[:space:]' || true)
  fi

  # Fallback: dpkg status for apt-installed libcuda1 (dir-type storage, no start needed)
  if [[ -z "$version" ]]; then
    local rootfs="/var/lib/lxc/${ctid}/rootfs"
    if [[ -f "${rootfs}/var/lib/dpkg/status" ]]; then
      version=$(grep -A5 "^Package: libcuda1$" "${rootfs}/var/lib/dpkg/status" \
        | grep "^Version:" | head -1 | awk '{print $2}' | cut -d- -f1)
    fi
  fi

  echo "${version:-$(translate 'not installed')}"
}

_detect_container_distro() {
  local distro
  distro=$(pct exec "$1" -- grep "^ID=" /etc/os-release 2>/dev/null \
    | cut -d= -f2 | tr -d '[:space:]"')
  echo "${distro:-unknown}"
}

_ensure_container_memory() {
  local ctid="$1"
  local cur_mem
  cur_mem=$(pct config "$ctid" 2>/dev/null | awk '/^memory:/{print $2}')
  [[ -z "$cur_mem" ]] && cur_mem=512

  if [[ "$cur_mem" -lt "$NVIDIA_INSTALL_MIN_MB" ]]; then
    if whiptail --title "$(translate 'Low Container Memory')" --yesno \
      "$(translate 'Container') ${ctid} $(translate 'has') ${cur_mem}MB RAM.\n\n$(translate 'The NVIDIA installer needs at least') ${NVIDIA_INSTALL_MIN_MB}MB $(translate 'to run without being killed by the OOM killer.')\n\n$(translate 'Increase container RAM temporarily to') ${NVIDIA_INSTALL_MIN_MB}MB?" \
      13 72; then
      CT_ORIG_MEM="$cur_mem"
      pct set "$ctid" -memory "$NVIDIA_INSTALL_MIN_MB" >>"$LOG_FILE" 2>&1 || true
    else
      msg_warn "$(translate 'Insufficient memory. Skipping LXC') ${ctid}."
      return 1
    fi
  fi
  return 0
}

_restore_container_memory() {
  local ctid="$1"
  if [[ -n "$CT_ORIG_MEM" ]]; then
    msg_info "$(translate 'Restoring container memory to') ${CT_ORIG_MEM}MB..."
    pct set "$ctid" -memory "$CT_ORIG_MEM" >>"$LOG_FILE" 2>&1 || true
    msg_ok "$(translate 'Memory restored.')"
    CT_ORIG_MEM=""
  fi
}

_start_container_and_wait() {
  local ctid="$1"
  msg_info "$(translate 'Starting container') ${ctid}..."
  pct start "$ctid" >>"$LOG_FILE" 2>&1 || true

  local ready=false
  for _ in {1..15}; do
    sleep 2
    if pct exec "$ctid" -- true >/dev/null 2>&1; then
      ready=true
      break
    fi
  done

  if ! $ready; then
    msg_warn "$(translate 'Container') ${ctid} $(translate 'did not become ready. Skipping.')"
    return 1
  fi
  msg_ok "$(translate 'Container') ${ctid} $(translate 'started.')" | tee -a "$screen_capture"
  return 0
}

update_lxc_nvidia() {
  local ctid="$1"
  local version="$2"
  local started_here=false

  local old_version
  old_version=$(get_lxc_nvidia_version "$ctid")

  msg_info2 "$(translate 'Container') ${ctid}: $(translate 'updating NVIDIA userspace libs') (${old_version} → ${version})"

  if ! pct status "$ctid" 2>/dev/null | grep -q "running"; then
    started_here=true
    _start_container_and_wait "$ctid" || return 1
  fi

  msg_info "$(translate 'Detecting container OS...')"
  local distro
  distro=$(_detect_container_distro "$ctid")
  msg_ok "$(translate 'Container OS:') ${distro}" | tee -a "$screen_capture"

  local install_rc=0

  case "$distro" in
    arch|manjaro|endeavouros)
      msg_info2 "$(translate 'Upgrading NVIDIA utils (Arch)...')"
      pct exec "$ctid" -- bash -c \
        "pacman -Syu --noconfirm nvidia-utils" \
        2>&1 | tee -a "$LOG_FILE"
      install_rc=${PIPESTATUS[0]}
      ;;
    *)
      local run_file="${NVIDIA_WORKDIR}/NVIDIA-Linux-x86_64-${version}.run"

      if [[ ! -f "$run_file" ]]; then
        msg_warn "$(translate 'Installer not found:') ${run_file}. $(translate 'Skipping LXC') ${ctid}."
        install_rc=1
      elif ! _ensure_container_memory "$ctid"; then
        install_rc=1
      else
        local free_mb
        free_mb=$(pct exec "$ctid" -- df -P -m / 2>/dev/null | awk 'END{print $4}')
        free_mb=${free_mb:-0}
        if [[ "$free_mb" -lt 1500 ]]; then
          _restore_container_memory "$ctid"
          whiptail --backtitle "ProxMenux" \
            --title "$(translate 'Insufficient Disk Space')" \
            --msgbox "\n$(translate 'Container') ${ctid} $(translate 'has only') ${free_mb}MB $(translate 'of free disk space.')\n\n$(translate 'NVIDIA libs require approximately 1.5GB of free space.')" \
            11 72
          msg_warn "$(translate 'Insufficient disk space. Skipping LXC') ${ctid}."
          install_rc=1
        else
          local extract_dir="${NVIDIA_WORKDIR}/extracted_${version}"
          local archive="/tmp/nvidia_lxc_${version}.tar.gz"

          msg_info2 "$(translate 'Extracting NVIDIA installer on host...')"
          rm -rf "$extract_dir"
          sh "$run_file" --extract-only --target "$extract_dir" 2>&1 | tee -a "$LOG_FILE"
          if [[ ${PIPESTATUS[0]} -ne 0 ]]; then
            msg_warn "$(translate 'Extraction failed. Check log:') ${LOG_FILE}"
            _restore_container_memory "$ctid"
            install_rc=1
          else
            msg_ok "$(translate 'NVIDIA installer extracted.')" | tee -a "$screen_capture"

            msg_info2 "$(translate 'Packing installer archive...')"
            tar --checkpoint=5000 --checkpoint-action=dot \
                -czf "$archive" -C "$extract_dir" . 2>&1 | tee -a "$LOG_FILE"
            echo ""
            local archive_size
            archive_size=$(du -sh "$archive" 2>/dev/null | cut -f1)
            msg_ok "$(translate 'Archive ready') (${archive_size})." | tee -a "$screen_capture"

            msg_info "$(translate 'Copying installer to container') ${ctid}..."
            if ! pct push "$ctid" "$archive" /tmp/nvidia_lxc.tar.gz >>"$LOG_FILE" 2>&1; then
              msg_warn "$(translate 'pct push failed. Check log:') ${LOG_FILE}"
              rm -f "$archive"
              rm -rf "$extract_dir"
              _restore_container_memory "$ctid"
              install_rc=1
            else
              rm -f "$archive"
              msg_ok "$(translate 'Installer copied to container.')" | tee -a "$screen_capture"

              msg_info2 "$(translate 'Running NVIDIA installer in container. This may take several minutes...')"
              echo "" >>"$LOG_FILE"
              if [[ "$distro" == "alpine" ]]; then
                # Alpine uses musl libc and does not ship a glibc dynamic
                # loader, so the nvidia-installer binary (glibc) cannot
                # execute. We pull `gcompat` to provide the glibc loader
                # and a libc shim, then copy the userspace libs and the
                # standard NVIDIA binaries by hand. SONAME symlinks are
                # built from `readelf` (binutils) instead of trusting a
                # hard-coded list — the .run ships ~50 .so files and the
                # set varies between branches.
                pct exec "$ctid" -- sh -c '
                  set -e
                  mkdir -p /tmp/nvidia_lxc_install
                  tar -xzf /tmp/nvidia_lxc.tar.gz -C /tmp/nvidia_lxc_install
                  apk add --no-cache gcompat binutils >/dev/null
                  cd /tmp/nvidia_lxc_install
                  mkdir -p /usr/lib /usr/bin
                  cp -P *.so* /usr/lib/ 2>/dev/null || true
                  for lib in /usr/lib/lib*.so.*; do
                    [ -f "$lib" ] || continue
                    soname=$(readelf -d "$lib" 2>/dev/null | grep SONAME | head -n1 | sed -e "s/.*\[//" -e "s/\].*//")
                    [ -n "$soname" ] && [ "$(basename "$lib")" != "$soname" ] && ln -sf "$(basename "$lib")" "/usr/lib/$soname"
                  done
                  for bin in nvidia-smi nvidia-debugdump nvidia-cuda-mps-control nvidia-cuda-mps-server nvidia-persistenced nvidia-modprobe; do
                    [ -f "$bin" ] && cp -P "$bin" /usr/bin/ && chmod 755 "/usr/bin/$bin"
                  done
                  rm -rf /tmp/nvidia_lxc_install /tmp/nvidia_lxc.tar.gz
                ' 2>&1 | tee -a "$LOG_FILE"
                install_rc=${PIPESTATUS[0]}
              else
                pct exec "$ctid" -- bash -c "
                  mkdir -p /tmp/nvidia_lxc_install
                  tar -xzf /tmp/nvidia_lxc.tar.gz -C /tmp/nvidia_lxc_install 2>&1
                  /tmp/nvidia_lxc_install/nvidia-installer \
                    --no-kernel-modules \
                    --no-questions \
                    --ui=none \
                    --no-nouveau-check \
                    --no-dkms \
                    --no-install-compat32-libs
                  EXIT=\$?
                  rm -rf /tmp/nvidia_lxc_install /tmp/nvidia_lxc.tar.gz
                  exit \$EXIT
                " 2>&1 | tee -a "$LOG_FILE"
                install_rc=${PIPESTATUS[0]}
              fi

              rm -rf "$extract_dir"
              _restore_container_memory "$ctid"
            fi
          fi
        fi
      fi
      ;;
  esac

  if [[ $install_rc -ne 0 ]]; then
    msg_warn "$(translate 'NVIDIA update failed for LXC') ${ctid} (rc=${install_rc}). $(translate 'Check log:') ${LOG_FILE}"
    if $started_here; then
      pct stop "$ctid" >>"$LOG_FILE" 2>&1 || true
    fi
    return 1
  fi

  if pct exec "$ctid" -- sh -c "which nvidia-smi" >/dev/null 2>&1; then
    local new_ver
    new_ver=$(pct exec "$ctid" -- nvidia-smi \
      --query-gpu=driver_version --format=csv,noheader 2>/dev/null \
      | head -1 | tr -d '[:space:]' || true)
    msg_ok "$(translate 'Container') ${ctid}: ${old_version} → ${new_ver:-$version}" | tee -a "$screen_capture"
  else
    msg_warn "$(translate 'nvidia-smi not found in container') ${ctid} $(translate 'after update.')"
  fi

  if $started_here; then
    msg_info "$(translate 'Stopping container') ${ctid}..."
    pct stop "$ctid" >>"$LOG_FILE" 2>&1 || true
    msg_ok "$(translate 'Container stopped.')" | tee -a "$screen_capture"
  fi
  return 0
}

# Post-host-install LXC update offer — scans for NVIDIA LXCs and, if any are
# found, asks the user if they want to propagate the driver update to them.
offer_lxc_updates_if_any() {
  local target_version="$1"
  find_nvidia_containers
  [[ ${#NVIDIA_CONTAINERS[@]} -eq 0 ]] && return 0

  local info ctid lxc_ver ct_name
  info="\n$(translate 'The following LXC containers have NVIDIA passthrough configured:')\n\n"
  for ctid in "${NVIDIA_CONTAINERS[@]}"; do
    lxc_ver=$(get_lxc_nvidia_version "$ctid")
    ct_name=$(pct config "$ctid" 2>/dev/null | grep "^hostname:" | awk '{print $2}')
    info+="  CT ${ctid}  ${ct_name:+(${ct_name})}  — $(translate 'driver:') ${lxc_ver}\n"
  done
  info+="\n$(translate 'Do you want to update the NVIDIA userspace libraries inside these containers to match the host?')"

  if ! hybrid_whiptail_yesno "$(translate 'Update NVIDIA in LXC Containers')" "$info" 20 80; then
    msg_info2 "$(translate 'LXC update skipped by user.')"
    return 0
  fi

  for ctid in "${NVIDIA_CONTAINERS[@]}"; do
    update_lxc_nvidia "$ctid" "$target_version" || true
  done
}

# ==========================================================
# System preparation (repos, headers, etc.)
# ==========================================================
ensure_repos_and_headers() {
  # Bootstrap APT repos FIRST. On a fresh Proxmox install the
  # pve-no-subscription / debian repos aren't configured by default
  # → `pve-headers-$(uname -r)` and `build-essential` come back as
  # "Unable to locate package" and the NVIDIA install bails out with
  # "no cc found". We delegate to the shared helper (same one the
  # post-install flow uses), which owns its own spinner pair — that's
  # why this block has to run BEFORE we open our own msg_info.
  if ! declare -F ensure_repositories >/dev/null 2>&1; then
    local _utils_install="$LOCAL_SCRIPTS/global/utils-install-functions.sh"
    [[ ! -f "$_utils_install" ]] && _utils_install="/usr/local/share/proxmenux/scripts/global/utils-install-functions.sh"
    # shellcheck source=/dev/null
    [[ -f "$_utils_install" ]] && source "$_utils_install"
  fi
  if declare -F ensure_repositories >/dev/null 2>&1; then
    ensure_repositories >>"$LOG_FILE" 2>&1 || true
  fi

  # Now own the spinner for the headers + build-tools check.
  msg_info "$(translate 'Checking kernel headers and build tools...')"

  local kver
  kver=$(uname -r)

  apt-get update -qq >>"$LOG_FILE" 2>&1

  if ! dpkg -s "pve-headers-$kver" >/dev/null 2>&1 && \
     ! dpkg -s "proxmox-headers-$kver" >/dev/null 2>&1; then
    apt-get install -y "pve-headers-$kver" "proxmox-headers-$kver" build-essential dkms >>"$LOG_FILE" 2>&1 || true
  else
    apt-get install -y build-essential dkms >>"$LOG_FILE" 2>&1 || true
  fi

  msg_ok "$(translate 'Kernel headers and build tools verified.')" | tee -a "$screen_capture"
}

blacklist_nouveau() {
  msg_info "$(translate 'Blacklisting nouveau driver...')"

  # Write blacklist config files
  if ! grep -q '^blacklist nouveau' /etc/modprobe.d/blacklist.conf 2>/dev/null; then
    echo "blacklist nouveau" >> /etc/modprobe.d/blacklist.conf
  fi

  # Also write explicit options file to ensure it's fully disabled
  cat > /etc/modprobe.d/nouveau-blacklist.conf <<'EOF'
blacklist nouveau
options nouveau modeset=0
EOF

  # Attempt to unload nouveau if currently loaded.
  # Close the spinner from the opening msg_info before going further —
  # otherwise the second msg_info below leaves the first one spinning
  # forever (visible on fresh installs where nouveau is loaded; invisible
  # on reinstalls where this branch is skipped).
  if grep -q "^nouveau " /proc/modules 2>/dev/null; then
    msg_ok "$(translate 'nouveau driver has been blacklisted.')" | tee -a "$screen_capture"

    msg_info "$(translate 'Nouveau module is loaded, attempting to unload...')"
    modprobe -r nouveau 2>/dev/null || true
    sleep 1

    # Check if unload succeeded
    if grep -q "^nouveau " /proc/modules 2>/dev/null; then
      NOUVEAU_STILL_LOADED=true
      msg_warn "$(translate 'Could not unload nouveau module (may be in use). The blacklist will take effect after reboot. Installation will continue but a reboot will be required.')"
      echo "WARNING: nouveau module still loaded after unload attempt" >> "$LOG_FILE"
    else
      NOUVEAU_STILL_LOADED=false
      msg_ok "$(translate 'nouveau module unloaded successfully.')" | tee -a "$screen_capture"
    fi
  else
    NOUVEAU_STILL_LOADED=false
    msg_ok "$(translate 'nouveau driver has been blacklisted.')" | tee -a "$screen_capture"
  fi

}

ensure_modules_config() {
  msg_info "$(translate 'Configuring NVIDIA modules...')"
  cat > /etc/modules-load.d/nvidia-vfio.conf <<'EOF'
nvidia
nvidia_uvm
EOF
  msg_ok "$(translate 'Modules configuration updated.')" | tee -a "$screen_capture"
}

stop_and_disable_nvidia_services() {
  local services=(
    "nvidia-persistenced.service"
    "nvidia-persistenced"
    "nvidia-powerd.service"
  )

  local services_detected=0 

  for service in "${services[@]}"; do
    if systemctl is-active --quiet "$service" 2>/dev/null || \
       systemctl is-enabled --quiet "$service" 2>/dev/null; then
      services_detected=1
      break 
    fi
  done

  if [ "$services_detected" -eq 1 ]; then
    msg_info "$(translate 'Stopping and disabling NVIDIA services...')"
    
    for service in "${services[@]}"; do
      if systemctl is-active --quiet "$service" 2>/dev/null; then
        systemctl stop "$service" >/dev/null 2>&1 || true
      fi
      if systemctl is-enabled --quiet "$service" 2>/dev/null; then
        systemctl disable "$service" >/dev/null 2>&1 || true
      fi
    done

    sleep 2
    
    msg_ok "$(translate 'NVIDIA services stopped and disabled.')" | tee -a "$screen_capture"
  fi
}

unload_nvidia_modules() {
  for mod in nvidia_uvm nvidia_drm nvidia_modeset nvidia; do
    modprobe -r "$mod" >/dev/null 2>&1 || true
  done

  # Give the kernel a moment to finalize sysfs teardown before re-checking.
  # Reading /proc/modules directly (instead of lsmod) avoids the
  # "could not open /sys/module/<mod>/holders" race when a module has just
  # been removed from /proc/modules but its sysfs dir hasn't been reaped yet.
  sleep 1

  if grep -q "^nvidia" /proc/modules 2>/dev/null; then
    for mod in nvidia_uvm nvidia_drm nvidia_modeset nvidia; do
      modprobe -r --force "$mod" >/dev/null 2>&1 || true
    done
    sleep 1
  fi

  if grep -q "^nvidia" /proc/modules 2>/dev/null; then

    if command -v lsof >/dev/null 2>&1; then
      echo "$(translate 'Processes using NVIDIA:'):" >> "$LOG_FILE"
      lsof /dev/nvidia* 2>/dev/null >> "$LOG_FILE" || true
    fi
  else

    msg_ok "$(translate 'NVIDIA kernel modules unloaded successfully.')" | tee -a "$screen_capture"
  fi
}

complete_nvidia_uninstall() {
  stop_and_disable_nvidia_services
  unload_nvidia_modules
  
  if command -v nvidia-uninstall >/dev/null 2>&1; then
    msg_info "$(translate 'Running NVIDIA uninstaller...')"
    nvidia-uninstall --silent >>"$LOG_FILE" 2>&1 || true
    msg_ok "$(translate 'NVIDIA uninstaller completed.')"
  fi
  msg_ok "$(translate 'NVIDIA uninstallation steps completed.')" | tee -a "$screen_capture"
  cleanup_nvidia_dkms
  
  msg_info "$(translate 'Removing NVIDIA packages...')"
  apt-get -y purge 'nvidia-*' 'libnvidia-*' 'cuda-*' 'libcudnn*' >>"$LOG_FILE" 2>&1 || true
  apt-get -y autoremove --purge >>"$LOG_FILE" 2>&1 || true
  apt-get -y autoclean >>"$LOG_FILE" 2>&1 || true
  
  rm -f /etc/modules-load.d/nvidia-vfio.conf
  rm -f /etc/udev/rules.d/70-nvidia.rules
  rm -rf /usr/lib/modprobe.d/nvidia*.conf
  rm -rf /etc/modprobe.d/nvidia*.conf
  
  if [[ -d "$NVIDIA_WORKDIR" ]]; then
    find "$NVIDIA_WORKDIR" -type d -name "nvidia-persistenced" -exec rm -rf {} + 2>/dev/null || true
    find "$NVIDIA_WORKDIR" -type d -name "nvidia-patch" -exec rm -rf {} + 2>/dev/null || true
  fi

  update_component_status "nvidia_driver" "removed" "" "gpu" '{}'
  
  msg_ok "$(translate 'Complete NVIDIA uninstallation finished.')" | tee -a "$screen_capture"
  
}

cleanup_nvidia_dkms() {
  local versions
  versions=$(dkms status 2>/dev/null | awk -F, '/nvidia/ {gsub(/ /,"",$2); print $2}' || true)

  [[ -z "$versions" ]] && return 0

  msg_info "$(translate 'Removing NVIDIA DKMS entries...')"
  while IFS= read -r ver; do
    [[ -z "$ver" ]] && continue
    dkms remove -m nvidia -v "$ver" --all >/dev/null 2>&1 || true
  done <<< "$versions"
  msg_ok "$(translate 'NVIDIA DKMS entries removed.')"
}

ensure_workdir() {
  mkdir -p "$NVIDIA_WORKDIR"
}

# ==========================================================
# Kernel + system detection
# ==========================================================
get_kernel_compatibility_info() {
  local kernel_version
  kernel_version=$(uname -r)

  if [[ -f /etc/pve/.version ]]; then
    PVE_VERSION=$(cat /etc/pve/.version)
  else
    PVE_VERSION="unknown"
  fi

  KERNEL_MAJOR=$(echo "$kernel_version" | cut -d. -f1)
  KERNEL_MINOR=$(echo "$kernel_version" | cut -d. -f2)

  MIN_DRIVER_VERSION=""
  RECOMMENDED_BRANCH=""
  COMPATIBILITY_NOTE=""
}

is_version_compatible() {
  return 0
}


is_current_nvidia_patched() {
  local status_file="/usr/local/share/proxmenux/components_status.json"
  [[ -f "$status_file" ]] || return 1
  command -v jq >/dev/null 2>&1 || return 1
  local patched
  patched=$(jq -r '.nvidia_driver.patched // false' "$status_file" 2>/dev/null)
  [[ "$patched" == "true" ]]
}

KEYLASE_PATCH_CACHE="/var/cache/proxmenux/keylase_patch_versions.txt"
KEYLASE_PATCH_TTL_SECONDS=$((7 * 86400))
KEYLASE_PATCH_URL="https://raw.githubusercontent.com/keylase/nvidia-patch/master/patch.sh"

# NVIDIA branch classification comes from the vendor's own Unix drivers
# page, not the CDN — the CDN publishes every branch (production, new
# feature, vulkan-beta, developer) in the same flat directory, whereas
# the vendor page carries the current heads clearly labelled "Production
# Branch", "New Feature Branch" and "Legacy GPU version". Extracting the
# majors from those three lines gives us the set of branches NVIDIA
# currently endorses for end users, with zero manual maintenance on our
# side — when NVIDIA promotes a new rama the cache picks it up on the
# next 24 h refresh. Cache is fail-open: if the fetch is blocked or the
# page layout changes, we skip the branch filter rather than emptying
# the picker.
NVIDIA_BRANCHES_CACHE="/var/cache/proxmenux/nvidia_stable_branches.txt"
NVIDIA_PRODUCTION_HEAD_CACHE="/var/cache/proxmenux/nvidia_production_head.txt"
NVIDIA_BRANCH_HEADS_CACHE="/var/cache/proxmenux/nvidia_branch_heads.txt"
NVIDIA_GPU_SUPPORT_CACHE_PREFIX="/var/cache/proxmenux/nvidia_gpu_support_"
NVIDIA_BRANCHES_TTL_SECONDS=$((24 * 3600))
NVIDIA_BRANCHES_URL="https://www.nvidia.com/en-us/drivers/unix/"

refresh_nvidia_branches_cache() {
  local now ts age
  now=$(date +%s)
  if [[ -f "$NVIDIA_BRANCHES_CACHE" ]]; then
    ts=$(stat -c '%Y' "$NVIDIA_BRANCHES_CACHE" 2>/dev/null || echo 0)
    age=$(( now - ts ))
    if (( age < NVIDIA_BRANCHES_TTL_SECONDS )) && [[ -s "$NVIDIA_BRANCHES_CACHE" ]]; then
      return 0
    fi
  fi
  mkdir -p "$(dirname "$NVIDIA_BRANCHES_CACHE")" 2>/dev/null || return 1
  local html tmp
  html=$(curl -fsSL -A "Mozilla/5.0" --max-time 15 "$NVIDIA_BRANCHES_URL" 2>/dev/null) || return 1
  [[ -z "$html" ]] && return 1
  local clean tmp_full
  clean=$(echo "$html" | perl -0777 -pe 's/<!--.*?-->//gs' 2>/dev/null)
  tmp=$(mktemp)
  tmp_full=$(mktemp)
  # Two label shapes on the page:
  #   • Production / New Feature → "… Branch Version:" then <a>full</a>.
  #   • Legacy                   → "Legacy GPU version (NNN.xx series):"
  #                                then <a>full</a> — the word "Version"
  #                                lives inside the parenthesised label
  #                                so the Production/Feature regex misses
  #                                it (separate alternative below).
  # HTML comments are stripped first so vestigial `<!-- Beta Version …
  # 387.34 -->` blocks in older ia32 rows don't leak stale majors.
  # Perl one-liner because we want to keep the label text next to the
  # version — shell greps can only give us one or the other, and we
  # need the label to tag each head as production / feature / legacy.
  echo "$clean" \
    | perl -ne '
        while (/(Production Branch Version|New Feature Branch Version|Legacy GPU version \(([0-9]+)\.xx series\)):[^<]*(?:<\/span>)?\s*<a[^>]*>([0-9]+\.[0-9]+(?:\.[0-9]+)?)/gi) {
          my $label = lc($1);
          my $ver   = $3;
          my ($maj) = split(/\./, $ver);
          my $type  = "unknown";
          if    ($label =~ /production branch/)  { $type = "production" }
          elsif ($label =~ /new feature branch/) { $type = "feature" }
          elsif ($label =~ /legacy gpu version/) { $type = "legacy" }
          print "$type|$maj|$ver\n";
        }' \
    | sort -u -t'|' -k1,1 -k2,2n > "$tmp_full"
  if [[ ! -s "$tmp_full" ]]; then
    rm -f "$tmp" "$tmp_full"
    return 1
  fi
  # Derive the majors-only file from the same source, but ONLY for
  # production + legacy — New Feature Branch heads (610.x today) are
  # intentionally NOT in the endorsed whitelist because they carry
  # kernel/driver features still under stabilisation. Superseded
  # production branches still qualify via the release-count heuristic
  # in filter_option_c_branch, so operators on 580 / 570 / 550 / 535
  # keep bugfix upgrade options.
  awk -F'|' '$1 == "production" || $1 == "legacy" { print $2 }' \
    "$tmp_full" | sort -un > "$tmp"
  if [[ ! -s "$tmp" ]]; then
    rm -f "$tmp" "$tmp_full"
    return 1
  fi
  mv "$tmp" "$NVIDIA_BRANCHES_CACHE"
  mv "$tmp_full" "$NVIDIA_BRANCH_HEADS_CACHE"

  # Extra pass: capture the full Production Branch head so the picker
  # can default to it instead of the highest numeric available (which
  # could be a New Feature Branch head — NVIDIA doesn't recommend those
  # as the general-purpose default). Best-effort.
  local prod_head
  prod_head=$(echo "$clean" \
    | grep -oiE 'Production Branch Version:[^<]*(</span>)?\s*<a[^>]*>[0-9]+\.[0-9]+(\.[0-9]+)?' \
    | grep -oE '>[0-9]+\.[0-9]+(\.[0-9]+)?' \
    | tr -d '>' \
    | head -n1)
  if [[ -n "$prod_head" ]]; then
    echo "$prod_head" > "$NVIDIA_PRODUCTION_HEAD_CACHE"
  else
    rm -f "$NVIDIA_PRODUCTION_HEAD_CACHE" 2>/dev/null || true
  fi
  return 0
}

# Return the full head version associated with a major from the
# branch-heads cache (e.g. `get_nvidia_branch_head 595` → 595.91.07).
# Cache lines are `type|major|version` — filter by major, print
# version. Used to know which release inside a branch to hit for the
# PCI-ID supported-GPUs list.
get_nvidia_branch_head() {
  local major="$1"
  [[ -f "$NVIDIA_BRANCH_HEADS_CACHE" && -s "$NVIDIA_BRANCH_HEADS_CACHE" ]] || return 1
  awk -F'|' -v m="$major" '$2 == m { print $3; exit }' "$NVIDIA_BRANCH_HEADS_CACHE"
}

# Refresh the per-branch supported-GPU cache. Uses the branch head as
# the "sample release" for the whole branch — NVIDIA rarely drops chip
# support inside a live branch, so this is a solid proxy that also
# minimises fetch count (~3 heads total instead of one per release).
# Written to nvidia_gpu_support_MAJOR.txt with one lowercase hex device
# id per line. Same 24h TTL as the branches cache.
refresh_nvidia_gpu_support_for_major() {
  local major="$1"
  [[ -z "$major" ]] && return 1
  local cache="${NVIDIA_GPU_SUPPORT_CACHE_PREFIX}${major}.txt"
  local now ts age
  now=$(date +%s)
  if [[ -f "$cache" ]]; then
    ts=$(stat -c '%Y' "$cache" 2>/dev/null || echo 0)
    age=$(( now - ts ))
    if (( age < NVIDIA_BRANCHES_TTL_SECONDS )) && [[ -s "$cache" ]]; then
      return 0
    fi
  fi
  local head_ver
  head_ver=$(get_nvidia_branch_head "$major") || return 1
  [[ -z "$head_ver" ]] && return 1
  local url="https://download.nvidia.com/XFree86/Linux-x86_64/${head_ver}/README/supportedchips.html"
  local html tmp
  html=$(curl -fsSL -A "Mozilla/5.0" --max-time 20 "$url" 2>/dev/null) || return 1
  [[ -z "$html" ]] && return 1
  mkdir -p "$(dirname "$cache")" 2>/dev/null || return 1
  tmp=$(mktemp)
  # NVIDIA's supportedchips.html lays out each GPU row as a <td> with
  # the PCI Device ID in 4-char hex. Anchor on the surrounding tag so
  # we don't sweep up unrelated 4-hex strings elsewhere in the page.
  echo "$html" \
    | grep -oiE '<td>[0-9A-F]{4}</td>' \
    | grep -oiE '[0-9A-F]{4}' \
    | tr 'A-F' 'a-f' \
    | sort -u > "$tmp"
  if [[ -s "$tmp" ]]; then
    mv "$tmp" "$cache"
    return 0
  fi
  rm -f "$tmp"
  return 1
}

# True if every detected NVIDIA GPU on this host has its device id in
# the branch's supported list. Fail-open when the cache is missing so
# a network hiccup never locks the picker out.
is_branch_compatible_with_host_gpus() {
  local major="$1"
  local cache="${NVIDIA_GPU_SUPPORT_CACHE_PREFIX}${major}.txt"
  [[ -f "$cache" && -s "$cache" ]] || return 0
  [[ ${#NVIDIA_HOST_GPU_IDS[@]} -eq 0 ]] && return 0
  local id
  for id in "${NVIDIA_HOST_GPU_IDS[@]}"; do
    grep -qFx "$id" "$cache" || return 1
  done
  return 0
}

# Reject Vulkan-beta / short-lived / developer branches by counting
# how many releases NVIDIA actually shipped inside that major on the
# CDN. Production and long-lived New Feature branches accumulate many
# release rows (470=17, 535=20, 550=14, 570=12, 580=14 …); Vulkan-beta
# and developer branches only ever get 1-4 releases before being
# superseded (590=2, 565=2, 530=2, 555=4 …). Threshold 5 separates the
# two groups cleanly at time of writing.
# Fail-open: if the release-count map hasn't been built for whatever
# reason, the branch passes (kernel + GPU-compat + curated whitelist
# are still enforced upstream). The map is populated once per
# `filter_option_c_branch` invocation, so no repeated CDN scraping.
NVIDIA_BRANCH_MIN_RELEASES=5
declare -A NVIDIA_BRANCH_RELEASE_COUNT=()

is_branch_release_count_sufficient() {
  local major="$1"
  [[ -z "$major" ]] && return 1
  [[ ${#NVIDIA_BRANCH_RELEASE_COUNT[@]} -eq 0 ]] && return 0
  local n="${NVIDIA_BRANCH_RELEASE_COUNT[$major]:-0}"
  (( n >= NVIDIA_BRANCH_MIN_RELEASES ))
}

get_nvidia_production_head() {
  [[ -f "$NVIDIA_PRODUCTION_HEAD_CACHE" && -s "$NVIDIA_PRODUCTION_HEAD_CACHE" ]] || return 1
  local v
  v=$(head -n1 "$NVIDIA_PRODUCTION_HEAD_CACHE" | tr -d '[:space:]')
  [[ -z "$v" ]] && return 1
  printf '%s\n' "$v"
}

is_nvidia_stable_branch() {
  local major="$1"
  [[ -z "$major" ]] && return 1
  # Fail-open: no cache → don't filter (upstream behaviour preserved).
  [[ -f "$NVIDIA_BRANCHES_CACHE" && -s "$NVIDIA_BRANCHES_CACHE" ]] || return 0
  grep -qFx "$major" "$NVIDIA_BRANCHES_CACHE"
}

refresh_keylase_patch_cache() {
  local now ts age
  now=$(date +%s)
  if [[ -f "$KEYLASE_PATCH_CACHE" ]]; then
    ts=$(stat -c '%Y' "$KEYLASE_PATCH_CACHE" 2>/dev/null || echo 0)
    age=$(( now - ts ))
    if (( age < KEYLASE_PATCH_TTL_SECONDS )) && [[ -s "$KEYLASE_PATCH_CACHE" ]]; then
      return 0
    fi
  fi
  mkdir -p "$(dirname "$KEYLASE_PATCH_CACHE")" 2>/dev/null || return 1
  local tmp
  tmp=$(mktemp)
  if curl -fsSL --max-time 15 "$KEYLASE_PATCH_URL" 2>/dev/null \
       | grep -oE '\["[0-9]+\.[0-9]+(\.[0-9]+)?"\]' \
       | sed -E 's/\["([0-9.]+)"\]/\1/' \
       | sort -u > "$tmp" && [[ -s "$tmp" ]]; then
    mv "$tmp" "$KEYLASE_PATCH_CACHE"
    return 0
  fi
  rm -f "$tmp"
  return 1
}

is_keylase_patch_supported() {
  local ver="$1"
  [[ -z "$ver" ]] && return 1
  [[ -f "$KEYLASE_PATCH_CACHE" && -s "$KEYLASE_PATCH_CACHE" ]] || return 1
  grep -qFx "$ver" "$KEYLASE_PATCH_CACHE"
}

filter_keylase_supported() {
  local versions_in="$1"
  while IFS= read -r ver; do
    [[ -z "$ver" ]] && continue
    if is_keylase_patch_supported "$ver"; then
      printf '%s\n' "$ver"
    fi
  done <<< "$versions_in"
}

filter_option_c_branch() {
  local versions_in="$1"
  local current="$2"
  local _unused_recommended_branch="$3"

  refresh_nvidia_branches_cache 2>/dev/null || true

  local target_branch=""
  if [[ -f "$NVIDIA_BRANCHES_CACHE" && -s "$NVIDIA_BRANCHES_CACHE" ]]; then
    target_branch=$(head -n1 "$NVIDIA_BRANCHES_CACHE")
  fi
  # Build a majors→count map from the incoming version list. This is
  # what backs `is_branch_release_count_sufficient` — done once per
  # call so the tight loop below stays local-arithmetic only.
  NVIDIA_BRANCH_RELEASE_COUNT=()
  while IFS= read -r _v; do
    [[ -z "$_v" ]] && continue
    local _m="${_v%%.*}"
    NVIDIA_BRANCH_RELEASE_COUNT[$_m]=$(( ${NVIDIA_BRANCH_RELEASE_COUNT[$_m]:-0} + 1 ))
  done <<< "$versions_in"
  # Grab the head (highest version) of every major so we know which
  # release to sample for supportedchips.html. We use the CDN listing
  # directly for this — the branch-heads cache only carries the
  # endorsed heads, not the superseded ones.
  declare -A _major_head=()
  while IFS= read -r _v; do
    [[ -z "$_v" ]] && continue
    local _m="${_v%%.*}"
    [[ -z "${_major_head[$_m]:-}" ]] && _major_head[$_m]="$_v"
  done < <(printf '%s\n' "$versions_in")
  # Warm the supported-GPU cache for every stable major (whitelist
  # heads: head already known → normal path; superseded heads: seed the
  # cache-file's head-version by directly writing a lightweight lookup).
  # For endorsed majors we can use refresh_nvidia_gpu_support_for_major
  # as-is (it looks up NVIDIA_BRANCH_HEADS_CACHE). For non-endorsed
  # majors we need to fetch supportedchips.html against the highest
  # release we saw in the CDN listing.
  local _m _head _cache _now _ts _age _html _tmp
  _now=$(date +%s)
  for _m in "${!_major_head[@]}"; do
    _head="${_major_head[$_m]}"
    _cache="${NVIDIA_GPU_SUPPORT_CACHE_PREFIX}${_m}.txt"
    if [[ -f "$_cache" ]]; then
      _ts=$(stat -c '%Y' "$_cache" 2>/dev/null || echo 0)
      _age=$(( _now - _ts ))
      if (( _age < NVIDIA_BRANCHES_TTL_SECONDS )) && [[ -s "$_cache" ]]; then
        continue
      fi
    fi
    _html=$(curl -fsSL -A "Mozilla/5.0" --max-time 20 \
      "https://download.nvidia.com/XFree86/Linux-x86_64/${_head}/README/supportedchips.html" \
      2>/dev/null) || continue
    [[ -z "$_html" ]] && continue
    mkdir -p "$(dirname "$_cache")" 2>/dev/null || continue
    _tmp=$(mktemp)
    echo "$_html" \
      | grep -oiE '<td>[0-9A-F]{4}</td>' \
      | grep -oiE '[0-9A-F]{4}' \
      | tr 'A-F' 'a-f' \
      | sort -u > "$_tmp"
    if [[ -s "$_tmp" ]]; then
      mv "$_tmp" "$_cache"
    else
      rm -f "$_tmp"
    fi
  done
  while IFS= read -r ver; do
    [[ -z "$ver" ]] && continue
    local ver_major="${ver%%.*}"
    if [[ -n "$target_branch" ]] && (( 10#$ver_major < 10#$target_branch )); then
      continue
    fi
    if is_nvidia_stable_branch "$ver_major" || is_branch_release_count_sufficient "$ver_major"; then
      if is_branch_compatible_with_host_gpus "$ver_major"; then
        printf '%s\n' "$ver"
      fi
    fi
  done <<< "$versions_in"
}

version_le() {
  local v1="$1"
  local v2="$2"

  IFS='.' read -r a1 b1 c1 <<<"$v1"
  IFS='.' read -r a2 b2 c2 <<<"$v2"

  a1=${a1:-0}; b1=${b1:-0}; c1=${c1:-0}
  a2=${a2:-0}; b2=${b2:-0}; c2=${c2:-0}

  a1=$((10#$a1)); b1=$((10#$b1)); c1=$((10#$c1))
  a2=$((10#$a2)); b2=$((10#$b2)); c2=$((10#$c2))

  if (( a1 < a2 )); then
    return 0
  elif (( a1 > a2 )); then
    return 1
  fi

  if (( b1 < b2 )); then
    return 0
  elif (( b1 > b2 )); then
    return 1
  fi

  if (( c1 <= c2 )); then
    return 0
  else
    return 1
  fi
}


# ==========================================================
# NVIDIA version management - FIXED VERSION
# ==========================================================
download_latest_version() {
  local latest_line version

  latest_line=$(curl -fsSL "${NVIDIA_BASE_URL}/latest.txt" 2>&1)
  if [[ -z "$latest_line" ]]; then
    echo "" >&2
    return 1
  fi

  version=$(echo "$latest_line" | awk '{print $1}' | tr -d '[:space:]')
  
  if [[ -z "$version" ]]; then
    echo "" >&2
    return 1
  fi
  
  if [[ ! "$version" =~ ^[0-9]+\.[0-9]+(\.[0-9]+)?$ ]]; then
    echo "" >&2
    return 1
  fi
  
  echo "$version"
  return 0
}

list_available_versions() {
  local html_content versions
  
  html_content=$(curl -s "$NVIDIA_BASE_URL/" 2>&1)
  
  if [[ -z "$html_content" ]]; then
    echo "" >&2
    return 1
  fi
  
  versions=$(echo "$html_content" \
    | grep -o 'href=[^ >]*' \
    | awk -F"'" '{print $2}' \
    | grep -E '^[0-9]' \
    | sed 's/\/$//' \
    | sed 's/^[[:space:]]*//;s/[[:space:]]*$//' \
    | sort -Vr \
    | uniq)
  
  if [[ -z "$versions" ]]; then
    echo "" >&2
    return 1
  fi
  
  echo "$versions"
  return 0
}

verify_version_exists() {
  local version="$1"
  local url="${NVIDIA_BASE_URL}/${version}/"
  
  if curl -fsSL --head "$url" >/dev/null 2>&1; then
    return 0
  else
    return 1
  fi
}

download_nvidia_installer() {
  ensure_workdir
  local version="$1"
  
  version=$(echo "$version" | tr -d '[:space:]' | tr -d '\n' | tr -d '\r')
  
  if [[ ! "$version" =~ ^[0-9]+\.[0-9]+(\.[0-9]+)?$ ]]; then
    msg_error "Invalid version format: $version" >&2
    echo "ERROR: Invalid version format: '$version'" >> "$LOG_FILE"
    return 1
  fi
  
  local run_file="$NVIDIA_WORKDIR/NVIDIA-Linux-x86_64-${version}.run"

  if [[ -f "$run_file" ]]; then
    echo "Found existing file: $run_file" >> "$LOG_FILE"
    local existing_size file_type
    existing_size=$(stat -c%s "$run_file" 2>/dev/null || stat -f%z "$run_file" 2>/dev/null || echo "0")
    file_type=$(file "$run_file" 2>/dev/null || echo "unknown")
    
    echo "Existing file size: $existing_size bytes" >> "$LOG_FILE"
    echo "Existing file type: $file_type" >> "$LOG_FILE"
    
 
    if [[ $existing_size -gt 40000000 ]] && echo "$file_type" | grep -q "executable"; then

      if sh "$run_file" --check 2>&1 | tee -a "$LOG_FILE" | grep -q "OK"; then
        echo "Existing file passed integrity check" >> "$LOG_FILE"
        msg_ok "$(translate 'Installer already downloaded and verified.')" >&2
        printf '%s\n' "$run_file"
        return 0
      else
        echo "Existing file FAILED integrity check, removing..." >> "$LOG_FILE"
        msg_warn "$(translate 'Existing file, re-downloading...')" >&2
        rm -f "$run_file"
      fi
    else
      echo "Existing file invalid (size or type), removing..." >> "$LOG_FILE"
      msg_warn "$(translate 'Removing invalid existing file...')" >&2
      rm -f "$run_file"
    fi
  fi

  if ! verify_version_exists "$version"; then
    msg_error "Version $version does not exist on NVIDIA servers" >&2
    echo "ERROR: Version $version not found on server" >> "$LOG_FILE"
    return 1
  fi

  local urls=(
    "${NVIDIA_BASE_URL}/${version}/NVIDIA-Linux-x86_64-${version}.run"
    "${NVIDIA_BASE_URL}/${version}/NVIDIA-Linux-x86_64-${version}-no-compat32.run"
  )

  # Web mode (ProxMenux Monitor) runs scripts without a controlling TTY, so
  # /dev/tty is not writable and progress-bar animations using \r don't render
  # in the web terminal. Fall back to a quiet wget in that case; interactive
  # users (SSH / console) still get the ISO-like progress bar.
  local _nv_has_tty=false
  if ! is_web_mode 2>/dev/null && [[ -t 2 ]]; then
    _nv_has_tty=true
  fi

  if $_nv_has_tty; then
    printf '\n  %s NVIDIA-Linux-x86_64-%s.run\n' \
      "$(translate 'Downloading')" "$version" >/dev/tty
  else
    echo "  $(translate 'Downloading') NVIDIA-Linux-x86_64-${version}.run" >&2
  fi

  local success=false
  local url_index=0

  for url in "${urls[@]}"; do
    ((url_index++))
    echo "Attempting download from: $url" >> "$LOG_FILE"

    rm -f "$run_file"

    local _dl_ok=false
    if $_nv_has_tty; then
      # Interactive: progress bar to /dev/tty (bypasses any caller redirection).
      if wget --no-verbose --show-progress \
              --connect-timeout=30 --timeout=600 --tries=1 \
              -O "$run_file" "$url" 2>/dev/tty; then
        _dl_ok=true
      fi
    else
      # Web / no-TTY: silent wget, log errors only.
      if wget --quiet \
              --connect-timeout=30 --timeout=600 --tries=1 \
              -O "$run_file" "$url" 2>>"$LOG_FILE"; then
        _dl_ok=true
      fi
    fi

    if $_dl_ok; then
      echo "Download completed, verifying file..." >> "$LOG_FILE"

      if [[ ! -f "$run_file" ]]; then
        echo "ERROR: File not created after download" >> "$LOG_FILE"
        continue
      fi

      local file_size
      file_size=$(stat -c%s "$run_file" 2>/dev/null || stat -f%z "$run_file" 2>/dev/null || echo "0")
      echo "Downloaded file size: $file_size bytes" >> "$LOG_FILE"

      if [[ $file_size -lt 40000000 ]]; then
        echo "ERROR: File too small ($file_size bytes, expected >40MB)" >> "$LOG_FILE"
        head -c 200 "$run_file" >> "$LOG_FILE" 2>&1
        rm -f "$run_file"
        continue
      fi

      local file_type
      file_type=$(file "$run_file" 2>/dev/null)
      echo "File type: $file_type" >> "$LOG_FILE"

      if echo "$file_type" | grep -q "executable"; then
        echo "SUCCESS: Valid executable downloaded" >> "$LOG_FILE"
        success=true
        break
      else
        echo "ERROR: Not a valid executable" >> "$LOG_FILE"
        head -c 200 "$run_file" | od -c >> "$LOG_FILE" 2>&1
        rm -f "$run_file"
      fi
    else
      echo "ERROR: wget failed for $url (exit code: $?)" >> "$LOG_FILE"
      rm -f "$run_file"
    fi
  done

  if ! $success; then
    msg_error "$(translate 'Download failed for all attempted URLs')" >&2
    msg_error "Version $version may not be available for your architecture" >&2
    echo "ERROR: All download attempts failed" >> "$LOG_FILE"
    return 1
  fi

  chmod +x "$run_file"
  echo "Installation file ready: $run_file" >> "$LOG_FILE"
  printf '%s\n' "$run_file"
}

# ==========================================================
# Installation / uninstallation
# ==========================================================
run_nvidia_installer() {
  local installer="$1"

  msg_info2 "$(translate 'Starting NVIDIA installer. This may take several minutes...')"
  echo "" >>"$LOG_FILE"
  echo "=== Running NVIDIA installer: $installer ===" >>"$LOG_FILE"

  # If nouveau is still loaded, rebuild initramfs first so the blacklist takes
  # effect for the installer sanity checks. Without this the .run installer
  # detects nouveau as active and aborts even when --disable-nouveau is passed.
  if [[ "${NOUVEAU_STILL_LOADED:-false}" == "true" ]]; then
    msg_info "$(translate 'Rebuilding initramfs to apply nouveau blacklist before installation...')"
    update-initramfs -u -k all >>"$LOG_FILE" 2>&1 || true
    proxmox-boot-tool refresh >>"$LOG_FILE" 2>&1 || true
    # Try one more time to unload nouveau after initramfs rebuild
    modprobe -r nouveau 2>/dev/null || true
    sleep 1
    if grep -q "^nouveau " /proc/modules 2>/dev/null; then
      echo "WARNING: nouveau still loaded after initramfs rebuild, proceeding with --no-nouveau-check" >> "$LOG_FILE"
      msg_warn "$(translate 'nouveau still active. Proceeding with installation. A reboot will be required for the driver to work.')"
    else
      NOUVEAU_STILL_LOADED=false
      msg_ok "$(translate 'nouveau module unloaded after initramfs rebuild.')" | tee -a "$screen_capture"
    fi
  fi

  local tmp_extract_dir="$NVIDIA_WORKDIR/tmp_extract"
  mkdir -p "$tmp_extract_dir"

  # --no-nouveau-check: prevents the installer from aborting when nouveau is
  # still loaded. The blacklist files are already in place; nouveau will be
  # gone after the reboot that the script offers at the end.
  sh "$installer" \
    --tmpdir="$tmp_extract_dir" \
    --no-questions \
    --ui=none \
    --disable-nouveau \
    --no-nouveau-check \
    --dkms \
    2>&1 | tee -a "$LOG_FILE"
  local rc=${PIPESTATUS[0]}
  echo "" >>"$LOG_FILE"
  
  rm -rf "$tmp_extract_dir"

  if [[ $rc -ne 0 ]]; then
    msg_error "$(translate 'NVIDIA installer reported an error. Check /tmp/nvidia_install.log')"
    update_component_status "nvidia_driver" "failed" "" "gpu" '{"patched":false}'
    return 1
  fi

  msg_ok "$(translate 'NVIDIA driver installed successfully.')" | tee -a "$screen_capture"
  return 0
}

remove_nvidia_driver() {
  complete_nvidia_uninstall
}

install_udev_rules_and_persistenced() {
  msg_info "$(translate 'Installing NVIDIA udev rules and persistence service...')"

  cat >/etc/udev/rules.d/70-nvidia.rules <<'EOF'
# /etc/udev/rules.d/70-nvidia.rules
KERNEL=="nvidia", RUN+="/bin/bash -c '/usr/bin/nvidia-smi -L'"
KERNEL=="nvidia_uvm", RUN+="/bin/bash -c '/usr/bin/nvidia-modprobe -c0 -u'"
EOF

  udevadm control --reload-rules
  udevadm trigger --subsystem-match=drm --subsystem-match=pci || true

  ensure_workdir
  cd "$NVIDIA_WORKDIR" || return 1
  # Pin to the last release tag so a hostile push to upstream `master`
  # can't slip arbitrary code into the install. Bump as needed; the
  # `--depth 1` keeps the clone fast. Audit Tier 6 — `nvidia-persistenced`
  # git clone sin pinning de versión.
  local NVIDIA_PERSISTENCED_TAG="${NVIDIA_PERSISTENCED_TAG:-575.64.05}"
  if [[ ! -d nvidia-persistenced ]]; then
    git clone --depth 1 --branch "$NVIDIA_PERSISTENCED_TAG" \
      https://github.com/NVIDIA/nvidia-persistenced.git >>"$LOG_FILE" 2>&1 \
      || git clone --depth 1 https://github.com/NVIDIA/nvidia-persistenced.git >>"$LOG_FILE" 2>&1 \
      || true
  fi

  if [[ -d nvidia-persistenced/init ]]; then
    cd nvidia-persistenced/init || return 1
    ./install.sh >>"$LOG_FILE" 2>&1 || true
  fi

  msg_ok "$(translate 'NVIDIA udev rules and persistence service installed.')" | tee -a "$screen_capture"
}

apply_nvidia_patch_if_needed() {
  # NVIDIA_PATCH_AUTO=yes|no skips the yes/no prompt for non-interactive
  # callers; unset preserves the interactive menu behavior.
  case "${NVIDIA_PATCH_AUTO:-}" in
    yes) : ;;
    no)
      msg_info2 "$(translate 'NVIDIA patch not applied.')"
      update_component_status "nvidia_driver" "installed" "$CURRENT_DRIVER_VERSION" "gpu" '{"patched":false}'
      return 0
      ;;
    *)
      if ! hybrid_whiptail_yesno "$(translate 'NVIDIA Patch')" \
        "\n$(translate 'Do you want to apply the optional NVIDIA patch to remove some GPU limitations?')"; then
        msg_info2 "$(translate 'NVIDIA patch not applied.')"
        update_component_status "nvidia_driver" "installed" "$CURRENT_DRIVER_VERSION" "gpu" '{"patched":false}'
        return 0
      fi
      ;;
  esac

  msg_info "$(translate 'Cloning and applying NVIDIA patch (keylase/nvidia-patch)...')"
  ensure_workdir
  cd "$NVIDIA_WORKDIR" || return 1
  # Pin keylase/nvidia-patch to a known-good commit. Override via env var
  # for forward-compat as new driver versions land. patch.sh ships a list
  # of supported drivers in the repo; if our running driver isn't covered
  # the patch silently no-ops, so we surface a warning before running.
  # Audit Tier 6 — `keylase/nvidia-patch` sin pinning + sin compat check.
  local NVIDIA_PATCH_REF="${NVIDIA_PATCH_REF:-master}"
  if [[ ! -d nvidia-patch ]]; then
    git clone --depth 1 --branch "$NVIDIA_PATCH_REF" \
      https://github.com/keylase/nvidia-patch.git >>"$LOG_FILE" 2>&1 \
      || git clone --depth 1 https://github.com/keylase/nvidia-patch.git >>"$LOG_FILE" 2>&1 \
      || true
  fi

  # Best-effort compatibility check: peek the supported-driver list in
  # patch.sh and warn if our driver isn't on it.
  if [[ -n "$CURRENT_DRIVER_VERSION" && -f nvidia-patch/patch.sh ]]; then
    if ! grep -qF "$CURRENT_DRIVER_VERSION" nvidia-patch/patch.sh 2>/dev/null; then
      msg_warn "$(translate 'NVIDIA driver') $CURRENT_DRIVER_VERSION $(translate 'is not in the patch.sh supported list. The patch may no-op or fail; review keylase/nvidia-patch README before continuing.')"
    fi
  fi

  if [[ -x nvidia-patch/patch.sh ]]; then
    cd nvidia-patch || return 1
    ./patch.sh >>"$LOG_FILE" 2>&1 || true
    msg_ok "$(translate 'NVIDIA patch applied - check README for supported versions.')"
    update_component_status "nvidia_driver" "installed" "$CURRENT_DRIVER_VERSION" "gpu" '{"patched":true}'
  else
    msg_warn "$(translate 'Could not run NVIDIA patch script. Please verify repository and driver version.')"
    update_component_status "nvidia_driver" "installed" "$CURRENT_DRIVER_VERSION" "gpu" '{"patched":false}'
  fi
}

restart_prompt() {
  if hybrid_whiptail_yesno "$(translate 'NVIDIA Drivers')" \
    "\n$(translate 'The installation/changes require a server restart to apply correctly. Do you want to reboot now?')"; then
    msg_success "$(translate 'Installation completed. Press Enter to continue...')"
    read -r
    msg_warn "$(translate 'Restarting the server...')"
    rm -f "$screen_capture"
    reboot
  else
    msg_success "$(translate 'Installation completed. Please reboot the server manually as soon as possible.')"
    msg_success "$(translate 'Completed. Press Enter to return to menu...')"
    read -r
    rm -f "$screen_capture"
  fi
}

# ==========================================================
# Dialog menus
# ==========================================================
show_action_menu_if_installed() {
  if ! $CURRENT_DRIVER_INSTALLED; then
    ACTION="install"
    return 0
  fi

  local menu_choices=(
    "install" "$(translate 'Reinstall/Update NVIDIA drivers')"
    "remove"  "$(translate 'Uninstall NVIDIA drivers and configuration')"
  )

  ACTION=$(hybrid_menu "ProxMenux" "$(translate 'NVIDIA Actions')\n\n$(translate 'Choose an action:')" 26 80 16 "${menu_choices[@]}") || ACTION="cancel"
}

show_install_overview() {
  local overview
  overview="\n$(translate 'This installation will:')\n\n"
  overview+=" • $(translate 'Install NVIDIA proprietary drivers')\n"
  overview+=" • $(translate 'Configure GPU passthrough with VFIO')\n"
  overview+=" • $(translate 'Blacklist nouveau driver')\n"
  overview+=" • $(translate 'Enable IOMMU support if not enabled')\n"
  overview+=" • $(translate 'Optionally update NVIDIA libs in LXC containers with passthrough')\n\n"

  overview+="$(translate 'Detected GPU(s):')\n"
  overview+="\Zb\Z4$DETECTED_GPUS_TEXT\Zn\n"

  overview+="\n\Zn$(translate 'Current status: ') "
  overview+="\Zb${CURRENT_STATUS_TEXT}\Zn\n"

  # Scan for LXC containers with NVIDIA passthrough and surface them in the
  # overview so the user knows upfront they will be offered a driver update.
  find_nvidia_containers
  if [[ ${#NVIDIA_CONTAINERS[@]} -gt 0 ]]; then
    overview+="\n$(translate 'LXC containers with NVIDIA passthrough:')\n"
    local ctid lxc_ver ct_name
    for ctid in "${NVIDIA_CONTAINERS[@]}"; do
      lxc_ver=$(get_lxc_nvidia_version "$ctid")
      ct_name=$(pct config "$ctid" 2>/dev/null | grep "^hostname:" | awk '{print $2}')
      overview+="  \Zb\Z4CT ${ctid}\Zn  ${ct_name:+(${ct_name})}  — $(translate 'driver:') ${lxc_ver}\n"
    done
  fi

  overview+="\n$(translate 'After confirming, you will be asked to choose the NVIDIA driver version to install.')\n\n"
  overview+="$(translate 'Do you want to continue?')"

  hybrid_yesno "$(translate 'NVIDIA GPU Driver Installation')" "$overview" 24 90
}

show_version_menu() {
  local latest versions_list
  local kernel_version
  kernel_version=$(uname -r)

  show_proxmenux_logo
  msg_title "$(translate 'NVIDIA GPU Driver Installation')"
  msg_info "$(translate 'Fetching compatible driver versions for your kernel and GPU...')"

  latest=$(download_latest_version 2>/dev/null)
  versions_list=$(list_available_versions 2>/dev/null)

  if [[ -z "$latest" ]] && [[ -z "$versions_list" ]]; then
    stop_spinner
    hybrid_msgbox "$(translate 'Error')" \
      "$(translate 'Could not retrieve versions list from NVIDIA. Please check your internet connection.')\n\nURL: ${NVIDIA_BASE_URL}" 10 80
    DRIVER_VERSION="cancel"
    return 1
  fi
  

  if [[ -z "$latest" ]] && [[ -n "$versions_list" ]]; then
    latest=$(echo "$versions_list" | head -n1)
  fi
  

  if [[ -n "$latest" ]] && [[ -z "$versions_list" ]]; then
    versions_list="$latest"
  fi
  
  # Clean latest version
  latest=$(echo "$latest" | tr -d '[:space:]')
  
  local current_list="$versions_list"
  
  # Apply kernel compatibility filter if needed
  if [[ -n "$MIN_DRIVER_VERSION" ]]; then
    local filtered_list=""
    while IFS= read -r ver; do
      [[ -z "$ver" ]] && continue
      if is_version_compatible "$ver"; then
        filtered_list+="$ver"$'\n'
      fi
    done <<< "$current_list"
    current_list="$filtered_list"
  fi

  if [[ -n "$current_list" ]]; then
    current_list=$(filter_option_c_branch "$current_list" "$CURRENT_DRIVER_VERSION" "")
  fi

  # Historically the picker capped candidates at `latest` (from the
  # CDN's `latest.txt`) so users never saw versions newer than the
  # global "latest". But latest.txt lags the Production Branch head
  # (595.91.07 today vs 595.84 in latest.txt) and also hides the New
  # Feature Branch head (610.x) that is a legitimate option once
  # kernel + GPU compat pass. Kernel floor, endorsement whitelist,
  # release-count heuristic and GPU-compat filter already narrow the
  # list to safe candidates; the Production head still stands out in
  # the "Latest available" recommendation, so an artificial ceiling
  # only masked valid options.

  # If the user has the keylase NVENC patch applied, only offer versions
  # that the patch supports — picking an unsupported version reinstalls
  # the driver fine but the patch silently no-ops afterwards, so the
  # user loses NVENC limit removal without warning.
  local patch_filtered=false
  local patch_filter_note=""
  if is_current_nvidia_patched && [[ -n "$current_list" ]]; then
    if refresh_keylase_patch_cache; then
      local trimmed
      trimmed=$(filter_keylase_supported "$current_list")
      if [[ -n "$trimmed" ]]; then
        current_list="$trimmed"
        patch_filtered=true
      else
        patch_filter_note="$(translate 'No version in this branch is currently supported by keylase/nvidia-patch — the NVENC patch will not reapply after reinstall.')"
      fi
    else
      patch_filter_note="$(translate 'Could not fetch keylase/nvidia-patch supported list — patch reapply compatibility is not verified.')"
    fi
  fi

  # Pick the default "Recommended" version. Three-tier priority so the
  # picker stays consistent with what the Monitor's driver-update
  # notification promised the user:
  #   1. If a driver is already installed AND its branch is still
  #      offered in the filtered list, recommend the highest release
  #      of that same branch (bugfix upgrade in place). Matches the
  #      Monitor's Hardware card, which surfaces "v580.178.04
  #      available" for a 580.x install — the user hitting Actualizar
  #      then expects to land on 580.178.04, not a cross-branch jump
  #      to Production. Cross-branch is still one row away in the
  #      list.
  #   2. Fresh install (no current driver) → Production Branch head
  #      from NVIDIA's Unix drivers page, when present in the list.
  #   3. Fallback → highest numeric in the list (Production may have
  #      been filtered out by kernel-compat / GPU-compat / patch
  #      awareness).
  latest=""
  if [[ -n "$CURRENT_DRIVER_VERSION" && -n "$current_list" ]]; then
    local _cur_branch="${CURRENT_DRIVER_VERSION%%.*}"
    if [[ -n "$_cur_branch" ]]; then
      local _same_branch_head
      _same_branch_head=$(printf '%s\n' "$current_list" \
        | awk -F. -v b="$_cur_branch" '$1 == b { print; exit }' \
        | tr -d '[:space:]')
      if [[ -n "$_same_branch_head" ]]; then
        latest="$_same_branch_head"
      fi
    fi
  fi
  if [[ -z "$latest" ]]; then
    local prod_head=""
    prod_head=$(get_nvidia_production_head 2>/dev/null) || prod_head=""
    if [[ -n "$prod_head" && -n "$current_list" ]]; then
      if printf '%s\n' "$current_list" | grep -qFx "$prod_head"; then
        latest="$prod_head"
      fi
    fi
  fi
  if [[ -z "$latest" && -n "$current_list" ]]; then
    latest=$(printf '%s\n' "$current_list" | head -n1 | tr -d '[:space:]')
  fi

  local menu_text="$(translate 'Select the NVIDIA driver version to install:')\n\n"
  menu_text+="$(translate 'Versions shown are compatible with your kernel and your GPU. The recommended version keeps you on your current driver branch, or defaults to the NVIDIA Production Branch head on a fresh install.')"
  if $patch_filtered; then
    menu_text+="\n\n$(translate 'NVENC patch detected — list narrowed to versions supported by keylase/nvidia-patch.')"
  elif [[ -n "$patch_filter_note" ]]; then
    menu_text+="\n\n${patch_filter_note}"
  fi

  local choices=()
  if [[ -n "$latest" ]]; then
    choices+=("$latest" "$latest — $(translate 'Recommended')")
  else
    choices+=("" "$(translate 'Recommended')")
  fi
  choices+=("" "")

  if [[ -n "$current_list" ]]; then
    while IFS= read -r ver; do
      [[ -z "$ver" ]] && continue
      ver=$(echo "$ver" | tr -d '[:space:]')
      [[ -z "$ver" ]] && continue
      
      choices+=("$ver" "$ver")
    done <<< "$current_list"
  else
    choices+=("" "$(translate 'No compatible versions found for your kernel')")
  fi

  stop_spinner
  local selection=$(hybrid_menu "$(translate 'NVIDIA Driver Version')" "$menu_text" 26 90 16 "${choices[@]}") || { DRIVER_VERSION="cancel"; return 1; }

  case "$selection" in
    "")
      DRIVER_VERSION="cancel"
      return 1
      ;;
    *)
      DRIVER_VERSION=$(echo "$selection" | tr -d '[:space:]')
      return 0
      ;;
  esac
}

# ==========================================================
# Main flow
# ==========================================================
main() {
  # Rotate the previous run's log instead of truncating — when the
  # current install fails, the user can compare against the previous
  # attempt to see what changed. Audit Tier 7 — log truncation.
  if [[ -f "$LOG_FILE" && -s "$LOG_FILE" ]]; then
    cp -p "$LOG_FILE" "${LOG_FILE}.prev" 2>/dev/null || true
  fi
  : >"$LOG_FILE"
  : >"$screen_capture"

  NOUVEAU_STILL_LOADED=false

  detect_nvidia_gpus
  detect_driver_status
  check_gpu_not_in_vm_passthrough
  check_stale_vfio_config_for_nvidia

  if ! $NVIDIA_GPU_PRESENT; then
    dialog --backtitle "ProxMenux" --title "$(translate 'NVIDIA GPU Driver Installation')" --msgbox \
      "\n$(translate 'No NVIDIA GPU has been detected on this system. The installer will now exit.')" 20 70
    exit 1
  fi

  show_action_menu_if_installed

  case "$ACTION" in
    install)
      if ! show_install_overview; then
        exit 0
      fi

      get_kernel_compatibility_info

      show_version_menu
      if [[ "$DRIVER_VERSION" == "cancel" || -z "$DRIVER_VERSION" ]]; then
        exit 0
      fi

      if $CURRENT_DRIVER_INSTALLED; then
        if [[ "$CURRENT_DRIVER_VERSION" == "$DRIVER_VERSION" ]]; then
          local confirm_text
          confirm_text="\n\n\n$(translate 'Version') \Zb\Z4$DRIVER_VERSION\Zn\n\n$(translate 'is already installed. Do you want to reinstall it? This will perform a clean uninstall first.')"
          if ! hybrid_yesno "$(translate 'Same Version Detected')" "$confirm_text" 14 70; then
              exit 0
          fi
        else
          local confirm_text
          confirm_text="\n\n$(translate 'Current version:') \Zb$CURRENT_DRIVER_VERSION\Zn\n"
          confirm_text+="$(translate 'New version:') \Zb\Z4$DRIVER_VERSION\Zn\n\n"
          confirm_text+="$(translate 'The current driver will be completely uninstalled before installing the new version. Continue?')"
          if ! hybrid_yesno "$(translate 'Version Change Detected')" "$confirm_text" 20 70; then
              exit 0
          fi
        fi
        
        show_proxmenux_logo
        msg_title "$(translate "$SCRIPT_TITLE")"
        msg_info2 "$(translate 'Uninstalling current NVIDIA driver before installing new version')"
        complete_nvidia_uninstall
        
        sleep 2
        
        CURRENT_DRIVER_INSTALLED=false
        CURRENT_DRIVER_VERSION=""
      fi

      show_proxmenux_logo
      msg_title "$(translate "$SCRIPT_TITLE")"

      ensure_repos_and_headers
      blacklist_nouveau
      ensure_modules_config
      
      stop_and_disable_nvidia_services
      unload_nvidia_modules

      local installer
      installer=$(download_nvidia_installer "$DRIVER_VERSION")
      local download_result=$?

      if [[ $download_result -ne 0 ]]; then
        msg_error "$(translate 'Failed to download NVIDIA installer')"
        exit 1
      fi

      msg_ok "$(translate 'NVIDIA installer downloaded successfully')" | tee -a "$screen_capture"

      if [[ -z "$installer" || ! -f "$installer" ]]; then
        msg_error "$(translate 'Internal error: NVIDIA installer path is empty or file not found.')"
        rm -f "$screen_capture"
        exit 1
      fi

      if ! run_nvidia_installer "$installer"; then
        rm -f "$screen_capture"
        exit 1
      fi
      
      sleep 2
      show_proxmenux_logo
      msg_title "$(translate "$SCRIPT_TITLE")"
      cat "$screen_capture"
      echo -e "${TAB}${GN}📄 $(translate "Log file")${CL}: ${BL}$LOG_FILE${CL}"

      install_udev_rules_and_persistenced

      msg_info "$(translate 'Updating initramfs for all kernels...')"
      update-initramfs -u -k all >>"$LOG_FILE" 2>&1 || true
      proxmox-boot-tool refresh >>"$LOG_FILE" 2>&1 || true
      msg_ok "$(translate 'initramfs updated.')"

      msg_info2 "$(translate 'Checking NVIDIA driver status with nvidia-smi')"
      if command -v nvidia-smi >/dev/null 2>&1; then
        nvidia-smi || true
        CURRENT_DRIVER_VERSION=$(nvidia-smi --query-gpu=driver_version --format=csv,noheader 2>/dev/null | head -n1)
        CURRENT_DRIVER_INSTALLED=true
      else
        msg_warn "$(translate 'nvidia-smi not found in PATH. Please verify the driver installation.')"
      fi

      if [[ -n "$CURRENT_DRIVER_VERSION" ]]; then
        msg_ok "$(translate 'NVIDIA driver') $CURRENT_DRIVER_VERSION $(translate 'installed successfully.')"
        update_component_status "nvidia_driver" "installed" "$CURRENT_DRIVER_VERSION" "gpu" '{"patched":false}'
        msg_success "$(translate 'Driver installed successfully. Press Enter to continue...')"
        read -r
      else
        msg_error "$(translate 'Failed to detect installed NVIDIA driver version.')"
        update_component_status "nvidia_driver" "failed" "" "gpu" '{"patched":false}'
      fi

      # Propagate the new driver to LXC containers with NVIDIA passthrough, if any.
      # Uses the same .run installer cached in $NVIDIA_WORKDIR — runs only if the
      # host install succeeded and the user confirms.
      if [[ -n "$CURRENT_DRIVER_VERSION" ]]; then
        offer_lxc_updates_if_any "$CURRENT_DRIVER_VERSION"
      fi

      apply_nvidia_patch_if_needed
      restart_prompt
      ;;
    remove)
      if hybrid_yesno "$(translate 'NVIDIA Driver Uninstall')" \
        "\n\n\n$(translate 'This will remove NVIDIA drivers and related configuration. Do you want to continue?')" 14 70; then

        show_proxmenux_logo
        msg_title "$(translate "$SCRIPT_TITLE")"

        remove_nvidia_driver

        msg_info "$(translate 'Updating initramfs for all kernels...')"
        update-initramfs -u -k all >>"$LOG_FILE" 2>&1 || true
        proxmox-boot-tool refresh >>"$LOG_FILE" 2>&1 || true
        msg_ok "$(translate 'initramfs updated.')"

        restart_prompt
      fi
      ;;
    cancel|*)
      exit 0
      ;;
  esac
}

# ==========================================================
# Non-interactive auto-reinstall entry point
# ==========================================================
# Invoked after a host-config restore by apply_cluster_postboot.sh
# when components_status.json reports nvidia_driver as installed
# but the kernel module isn't loaded on the live system (i.e. the
# restore brought back the configs but not the binary driver from
# /lib/modules/<kernel>/). Replays the install path the user
# originally ran via `menu → 2`, using the recorded version, with
# no dialogs.
#
# Exit codes:
#   0  installed (or no-op — GPU absent / driver already present)
#   1  state file unreadable or no nvidia_driver entry
#   2  install failed
auto_reinstall_from_state() {
  : >"$LOG_FILE"
  echo "=== auto_reinstall_from_state started $(date -Iseconds) ===" >>"$LOG_FILE"

  if ! command -v jq >/dev/null 2>&1; then
    echo "jq not available — cannot read components_status.json" | tee -a "$LOG_FILE"
    return 1
  fi
  if [[ ! -f "$COMPONENTS_STATUS_FILE" ]]; then
    echo "No components_status.json at $COMPONENTS_STATUS_FILE" | tee -a "$LOG_FILE"
    return 1
  fi

  local recorded_status recorded_version recorded_patched
  recorded_status=$(jq -r '.nvidia_driver.status // ""'    "$COMPONENTS_STATUS_FILE" 2>/dev/null)
  recorded_version=$(jq -r '.nvidia_driver.version // ""'  "$COMPONENTS_STATUS_FILE" 2>/dev/null)
  recorded_patched=$(jq -r '.nvidia_driver.patched // false' "$COMPONENTS_STATUS_FILE" 2>/dev/null)

  if [[ "$recorded_status" != "installed" ]]; then
    echo "nvidia_driver not marked installed in state ($recorded_status) — nothing to do" | tee -a "$LOG_FILE"
    return 0
  fi
  if [[ -z "$recorded_version" || "$recorded_version" == "null" ]]; then
    echo "nvidia_driver marked installed but no version recorded — aborting" | tee -a "$LOG_FILE"
    return 1
  fi
  echo "Recorded driver: $recorded_version (patched=$recorded_patched)" >>"$LOG_FILE"

  detect_nvidia_gpus
  if ! $NVIDIA_GPU_PRESENT; then
    echo "No NVIDIA GPU detected on this host — skipping reinstall" | tee -a "$LOG_FILE"
    return 0
  fi
  # Skip when the host is configured for VFIO passthrough. Two
  # signals — config wins over runtime:
  #
  #   1. modprobe.d declares the nvidia* modules blacklisted (the
  #      file ProxMenux drops when switching the GPU to VM mode is
  #      `proxmenux-nvidia-vfio-blacklist.conf`, but any blacklist
  #      file counts).
  #   2. The PCI device is bound to vfio-pci.
  #
  # The config check has to come FIRST because right after a host
  # restore + reboot, the binding may not yet be effective (the
  # GPU temporarily shows "no driver" while udev/initramfs settle).
  # Reinstalling NVIDIA in that window forces the module onto a
  # PCI device vfio-pci has already reserved → "NVRM: Try unloading
  # the conflicting kernel module" and exit 1.
  if [[ -f /etc/modprobe.d/proxmenux-nvidia-vfio-blacklist.conf ]] \
     || grep -qrhE '^[[:space:]]*blacklist[[:space:]]+nvidia([[:space:]_]|$)' \
          /etc/modprobe.d/ 2>/dev/null; then
    echo "NVIDIA blacklisted (host is VFIO-configured) — skipping host driver reinstall" \
      | tee -a "$LOG_FILE"
    return 0
  fi
  local _vfio_dev
  for _vfio_dev in /sys/bus/pci/devices/*; do
    [[ "$(cat "$_vfio_dev/vendor" 2>/dev/null)" == "0x10de" ]] || continue
    if [[ -L "$_vfio_dev/driver" ]] \
       && [[ "$(basename "$(readlink "$_vfio_dev/driver")")" == "vfio-pci" ]]; then
      echo "NVIDIA GPU bound to vfio-pci — skipping host driver reinstall" | tee -a "$LOG_FILE"
      return 0
    fi
  done
  detect_driver_status
  if $CURRENT_DRIVER_INSTALLED && [[ "$CURRENT_DRIVER_VERSION" == "$recorded_version" ]]; then
    echo "Driver $recorded_version already installed and matches state — no-op" | tee -a "$LOG_FILE"
    return 0
  fi

  DRIVER_VERSION="$recorded_version"

  # Same install path as the interactive main() flow, minus all
  # dialogs and confirmations.
  echo "Reinstalling NVIDIA driver $DRIVER_VERSION non-interactively..." | tee -a "$LOG_FILE"
  ensure_workdir
  ensure_repos_and_headers          >>"$LOG_FILE" 2>&1
  blacklist_nouveau                 >>"$LOG_FILE" 2>&1
  ensure_modules_config             >>"$LOG_FILE" 2>&1

  if $CURRENT_DRIVER_INSTALLED; then
    echo "Different version currently installed; cleaning up first..." | tee -a "$LOG_FILE"
    complete_nvidia_uninstall       >>"$LOG_FILE" 2>&1
  fi

  local installer
  installer=$(download_nvidia_installer "$DRIVER_VERSION" 2>>"$LOG_FILE")
  if [[ -z "$installer" || ! -f "$installer" ]]; then
    echo "Download failed — see $LOG_FILE" | tee -a "$LOG_FILE"
    return 2
  fi
  echo "Installer ready: $installer" >>"$LOG_FILE"
  if ! run_nvidia_installer "$installer" >>"$LOG_FILE" 2>&1; then
    echo "Install failed — see $LOG_FILE" | tee -a "$LOG_FILE"
    return 2
  fi
  install_udev_rules_and_persistenced >>"$LOG_FILE" 2>&1

  # Preserve the recorded patched state across the reinstall. The patch
  # helper writes its own update_component_status on success.
  CURRENT_DRIVER_VERSION="$DRIVER_VERSION"
  if [[ "$recorded_patched" == "true" ]]; then
    echo "Recorded state had patched=true — re-applying NVIDIA patch..." | tee -a "$LOG_FILE"
    NVIDIA_PATCH_AUTO=yes apply_nvidia_patch_if_needed >>"$LOG_FILE" 2>&1 || true
  else
    if declare -F update_component_status >/dev/null 2>&1; then
      update_component_status "nvidia_driver" "installed" "$DRIVER_VERSION" "gpu" '{"patched":false}' >>"$LOG_FILE" 2>&1
    fi
  fi

  echo "✓ NVIDIA driver $DRIVER_VERSION reinstalled" | tee -a "$LOG_FILE"
  return 0
}

if [[ "${BASH_SOURCE[0]}" == "$0" ]]; then
  if [[ "${1:-}" == "--auto-reinstall" ]]; then
    auto_reinstall_from_state
    exit $?
  fi
  main
fi