#!/bin/bash

PACKAGES="libcuda1 \
libegl-nvidia0 \
libgl1-nvidia-glvnd-glx \
libgles-nvidia1 \
libgles-nvidia2 \
libglx-nvidia0 \
libnvidia-cfg1 \
libnvidia-eglcore \
libnvidia-glcore \
libnvidia-glvkspirv \
libnvidia-ml1 \
libnvidia-nvvm4 \
libnvidia-ptxjitcompiler1 \
libnvidia-rtcore \
libxnvctrl0 \
nvidia-alternative \
nvidia-driver \
nvidia-driver-bin \
nvidia-driver-libs \
nvidia-egl-icd \
nvidia-kernel-dkms \
nvidia-kernel-support \
nvidia-settings \
nvidia-smi \
nvidia-vdpau-driver \
nvidia-vulkan-icd \
xserver-xorg-video-nvidia \
nvidia-persistenced"

function check_debian {
  if [[ -a /etc/debian_version ]]; then
    DEBIAN_VERSION=$(cat /etc/debian_version)
    if [[ ${DEBIAN_VERSION/.*} -ge 11 ]]; then
      return 0
    else
      echo -e "Not in Debian GNU/Linux version >= 11\nExit.\n"
      return 1
    fi
  else
    echo -e "Not in Debian GNU/Linux\nExit.\n"
    return 1
  fi
}

function show_versions {
  VERSIONS=($(apt-cache show nvidia-driver | grep -E "^Version" | awk '{print $2}'))
  echo -e "\n-= NVIDIA driver for Debian GNU/Linux v1.0 =-\n"
  for ((i = 0; i < ${#VERSIONS[@]}; i++)); do
    echo "$(( $i + 1 )). ${VERSIONS[$i]}"
  done
  echo -e "0. Exit.\n"
}

if [[ check_debian -eq 0 ]]; then
  dpkg -V cuda-keyring
  if [[ $? -eq 0 ]] && [[ $(id -u) -eq 0 ]]; then
    show_versions
    read -p "Select version: " CHOICE
    V=${VERSIONS[$((CHOICE-1))]}
    eval "case \"$CHOICE\" in $(seq -s '|' 1 ${#VERSIONS[@]}))
        echo -e \"Selected $V NVIDIA Driver version\n\"
        apt-mark unhold $PACKAGES && \
        apt update
        apt install -y nvidia-driver=$V \
        nvidia-driver-libs=$V \
        nvidia-driver-bin=$V \
        nvidia-vdpau-driver=$V \
        nvidia-kernel-dkms=$V \
        libgl1-nvidia-glvnd-glx=$V \
        nvidia-egl-icd=$V \
        libglx-nvidia0=$V \
        xserver-xorg-video-nvidia=$V \
        libnvidia-glcore=$V \
        libegl-nvidia0=$V \
        libnvidia-eglcore=$V \
        nvidia-settings=$V \
        libegl-nvidia0=$V \
        libglx-nvidia0=$V \
        libnvidia-cfg1=$V \
        nvidia-persistenced=$V \
        libxnvctrl0=$V \
        nvidia-kernel-support=$V \
        libgles-nvidia1=$V \
        libgles-nvidia2=$V \
        nvidia-vulkan-icd=$V \
        libnvidia-ml1=$V \
        libnvidia-glvkspirv=$V \
        libnvidia-rtcore=$V \
        libcuda1=$V \
        libnvidia-ptxjitcompiler1=$V \
        libnvidia-nvvm4=$V \
        nvidia-alternative=$V \
        nvidia-smi=$V
        if [[ \$? -eq 0 ]]; then
          apt clean && apt-mark hold $PACKAGES
          echo -e \"\nNVIDIA Driver $V installed sucessfully. Reboot.\n\"
          exit 0
        else
          apt clean && apt purge -y $PACKAGES
          echo -e \"\nNVIDIA Driver $V NOT installed. System clean.\n\"
          exit 1
        fi
        ;;
    0)
        echo -e \"\nExit. Byebye.\n\"
        exit 0
        ;;
    *)
        echo -e \"\nError: invalid NVIDIA driver version\n\"
        exit 1
        ;;
    esac"
  elif [[ $(id -u) -ne 0 ]]; then
    echo -e "\nThis script need root.\n$ sudo $0\nExit.\n"
    exit 0
  else
    wget -V
    if [[ $? -ne 0 ]]; then
      echo -e "\nInstalling wget...\n"
      apt install -y wget
    fi
    echo -e "\nInstalling NVIDIA keyring...\n"
    TMP="/tmp/cuda-keyring.deb"
    DEBIAN_VERSION=$(cat /etc/debian_version)
    wget -O$TMP "https://developer.download.nvidia.com/compute/cuda/repos/debian${DEBIAN_VERSION/.*}/x86_64/cuda-keyring_1.0-1_all.deb" && \
    dpkg -i $TMP && rm $TMP && apt update
    echo -e "\nExecuting script again...\n"
    exec $0
  fi
fi