# This script has to be executed manually as root with reboots in between

if [ -z "${1}" ]; then
  LAT_VERSIONS="$(wget -qO- https://www.nvidia.com/en-us/drivers/unix/ | tidy -quiet -wrap 4096 2>/dev/null | grep -A8 "Linux x86_64/AMD64/EM64T" | grep -e "Latest Production Branch" -e "Latest New Feature Branch" -e "Latest Beta Version")"
  NV_LAT="$(echo "$LAT_VERSIONS" | grep "Latest Production Branch" | cut -d '>' -f4 | cut -d '<' -f1)"
  NV_NFB="$(echo "$LAT_VERSIONS" | grep "Latest New Feature Branch" | cut -d '>' -f4 | cut -d '<' -f1)"
  NV_BETA="$(echo "$LAT_VERSIONS" | grep "Latest Beta Version" | cut -d '>' -f4 | cut -d '<' -f1)"
  echo -e "Available versions:\n---------------------------------\n"
  echo -e "Production Branch (prb):  ${NV_LAT}\nNew Feature Branch (nfb): ${NV_NFB}\nBeta Branch (beta):       ${NV_BETA}"
  echo -e "\nPlease use the sciprt with: ./nvidia.sh <BRANCH>\n                            ./nvidia.sh nfb"
  exit 0
fi

if [ ! -d ~/nvidia ]; then
  # Make sure that i386 architecture is available
  dpkg --add-architecture i386
  # Set Nvidia Driver version
  apt-get update
  apt-get -y upgrade
  apt-get -y autoremove
  apt-get -y install linux-headers-$(uname -r) build-essential libglvnd-dev pkg-config tidy
  mkdir -p ~/nvidia && cd ~/nvidia
  LAT_VERSIONS="$(wget -qO- https://www.nvidia.com/en-us/drivers/unix/ | tidy -quiet -wrap 4096 2>/dev/null | grep -A8 "Linux x86_64/AMD64/EM64T" | grep -e "Latest Production Branch" -e "Latest New Feature Branch" -e "Latest Beta Version")"
  if [ -z "${LAT_VERSIONS}" ]; then
    echo "ERROR: Something went wrong while grabbing the latest driver versions! Abort!"
    rm -rf ~/nvidia
    exit 1
  fi
  if [ "$1" == "prb" ]; then
    NV_DRV_V="$(echo "$LAT_VERSIONS" | grep "Latest Production Branch" | cut -d '>' -f4 | cut -d '<' -f1)"
  elif [ "$1" == "nfb" ]; then
    NV_DRV_V="$(echo "$LAT_VERSIONS" | grep "Latest New Feature Branch" | cut -d '>' -f4 | cut -d '<' -f1)"
  elif [ "$1" == "beta" ]; then
    NV_DRV_V="$(echo "$LAT_VERSIONS" | grep "Latest Beta Version" | cut -d '>' -f4 | cut -d '<' -f1)"
  else
    echo "ERROR: \"$1\" is not a valid option, valid options are:"
    echo
    echo -e "prb  = Latest Production branch\nnfb  = Latest New Feature branch\nbeta = Latest Beta branch"
    rm -rf ~/nvidia
    exit 1
  fi
  if [ -z "${NV_DRV_V}" ]; then
    echo "Error: Something went wrong while fetching driver version!"
    rm -rf ~/nvidia
    exit 1
  fi
  wget http://download.nvidia.com/XFree86/Linux-x86_64/${NV_DRV_V}/NVIDIA-Linux-x86_64-${NV_DRV_V}.run
  EXIT_STATUS=$?
  if [ "${EXIT_STATUS}" != "0" ]; then
    echo "ERROR: something went wrong with the download! Abort!"
    rm -rf ~/nvidia
    exit 1
  fi
  chmod +x NVIDIA-Linux-x86_64-${NV_DRV_V}.run
  echo "blacklist nouveau
  options nouveau modeset=0" > /etc/modprobe.d/blacklist-nouveau.conf
  export PATH=$PATH:/usr/bin:/usr/sbin
  update-initramfs -u
  systemctl set-default multi-user.target
  reboot now
fi

# After the reboot you are in the command line, login as root
apt-get update
apt-get -y upgrade
apt-get -y autoremove
apt-get -y install linux-headers-$(uname -r) build-essential libglvnd-dev pkg-config tidy
echo "Installing Nvidia driver, this will take some time..."
cd ~/nvidia
./NVIDIA-Linux-x86_64-*.run \
  --no-precompiled-interface \
  --disable-nouveau \
  --install-compat32-libs \
  --silent

# At the installer select "Yes" at "Install 32-bit compatibility libraries"
systemctl set-default graphical.target
rm -rf ~/nvidia

# If running a Secure Boot environment then you have to enroll the key:
# mokutil --import /usr/share/nvidia/<FILENAME>.der

reboot now