#!/bin/env bash

#    Compile latest Nvidia drivers on a Debian-based Linux
#    Copyright (C) 2018  Pekka Helenius
#
#    This program is free software: you can redistribute it and/or modify
#    it under the terms of the GNU General Public License as published by
#    the Free Software Foundation, either version 3 of the License, or
#    (at your option) any later version.
#
#    This program is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU General Public License for more details.
#
#    You should have received a copy of the GNU General Public License
#    along with this program.  If not, see <https://www.gnu.org/licenses/>.

##############################################################################

# Check if we're using bash or sh to run the script. If bash, OK.
# If another one, ask user to run the script with bash.

BASH_CHECK=$(ps | grep `echo $$` | awk '{ print $4 }')

if [ $BASH_CHECK != "bash" ]; then
    echo  "
Please run this script using bash (/usr/bin/bash).
    "
    exit 1
fi

########################################################

# Just a title & author for this script, used in initialization and help page

SCRIPT_TITLE="\e[1mNvidia drivers package builder & installer\e[0m"
SCRIPT_AUTHOR="Pekka Helenius (~Fincer), 2018"

########################################################

BUILD_MAINDIR=${PWD}/debian_nvidia

########################################################

_pkgname="nvidia"
arch="x86_64"
pkgver=396.54

files=(
"http://archive.ubuntu.com/ubuntu/pool/restricted/n/nvidia-graphics-drivers-390/nvidia-graphics-drivers-390_390.48-0ubuntu3.debian.tar.xz"

"http://us.download.nvidia.com/XFree86/Linux-${arch}/${pkgver}/NVIDIA-Linux-${arch}-${pkgver}.run"
)

###################

pkgver_major=$(printf '%s' ${pkgver} | grep -oE "^[0-9]+[^.]")
pkgdir="nvidia-graphics-drivers-${pkgver_major}_${pkgver}"

typeset -A library_fixes

###################

# From time to time, bundled library version numbers change
# in Nvidia driver packages. Update library version if needed
#
# Left side: old library version
# Right side: new library version
#
library_fixes=(
[libnvidia-egl-wayland.so.1.0.2]="libnvidia-egl-wayland.so.1.0.3"
)

###################

# These are defined build dependencies in debian/control file
nvidia_builddeps=(
'dpkg-dev'
'xz-utils'
'dkms'
'libwayland-client0'
'libwayland-server0'
'libxext6'
'quilt'
'po-debconf'
'execstack'
'dh-modaliases'
'xserver-xorg-dev'
'libglvnd-dev'
)

###################

# These packages are required by the compiled Nvidia packages
nvidia_required_packages=(

# Required by libnvidia-gl
"libwayland-client0"
"libwayland-server0"

# Required by libnvidia, libnvidia-decode & libnvidia-fbc1
"libx11-6"

# Required by libnvidia, libnvidia-decode, libnvidia-fbc1 & libnvidia-ifr1
"libxext6"

# Required by libnvidia-fbc1 & libnvidia-ifr1
"libgl1"

# Required by xserver-xorg-video-nvidia
"xserver-xorg-core"
"xorg-video-abi-23"

# Required by nvidia-compute-utils
"adduser"

)

###################

# Nvidia packages. THIS ORDER IS MANDATORY, DO NOT CHANGE!
nvidia_install_packages=(
# Similar than 'nvidia-dkms' package on Arch Linux
"nvidia-kernel-source-${pkgver_major}"

# Nvidia DKMS
"nvidia-kernel-common-${pkgver_major}"
"nvidia-dkms-${pkgver_major}"

# Similar than 'nvidia-utils' package on Arch Linux
"libnvidia-common-${pkgver_major}"
"libnvidia-gl-${pkgver_major}"
"libnvidia-cfg1-${pkgver_major}"
"xserver-xorg-video-nvidia-${pkgver_major}"
"libnvidia-compute-${pkgver_major}"
"libnvidia-decode-${pkgver_major}"
"libnvidia-encode-${pkgver_major}"
"libnvidia-fbc1-${pkgver_major}"
"libnvidia-ifr1-${pkgver_major}"
"nvidia-compute-utils-${pkgver_major}"
"nvidia-utils-${pkgver_major}"
)

########################################################

pkgver_sed=$(printf '%s' ${pkgver} | sed 's/\./\\\./')

i=0
for f in ${files[@]}; do
  file_basename=$(printf '%s' ${f} | awk -F / '{print $NF}')
  filebases[$i]=${file_basename}
  let i++
done

oldver=$(printf '%s' ${filebases[0]} | grep -oE "[0-9]{3}\.[0-9]{2}" | head -1)
oldver_major=$(printf '%s' ${filebases[0]} | grep -oE "[0-9]{3}" | head -1)
oldver_sed=$(printf '%s' ${oldver} | sed 's/\./\\\./')

########################################################

# http://wiki.bash-hackers.org/snipplets/print_horizontal_line#a_line_across_the_entire_width_of_the_terminal
function INFO_SEP() { printf '%*s\n' "${COLUMNS:-$(tput cols)}" '' | tr ' ' - ; }

###########################################################

echo -e "\n${SCRIPT_TITLE}\n"

########################################################

echo -e "Selected driver version:\t${pkgver}\n"
INFO_SEP

########################################################

mkdir -p ${BUILD_MAINDIR}

if [[ $? -eq 0 ]]; then
  cd ${BUILD_MAINDIR}
  WORKDIR=${PWD}
else
  echo -e "Error: couldn't create Nvidia build directory. Aborting\n"
  exit 1
fi

########################################################

# If the script is interrupted (Ctrl+C/SIGINT), do the following

function Nvidia_intCleanup() {
  rm -rf ${WORKDIR}/{${pkgdir}/{debian,NVIDIA-Linux,NVIDIA-Linux-amd64,NVIDIA-Linux-x86_64-${pkgver},LICENSE.txt,unpack-stamp},compiled_deb,compiled_other}
  exit 0
}

# Allow interruption of the script at any time (Ctrl + C)
trap "Nvidia_intCleanup" INT

###########################################################

# Remove old build files
rm -rf ${WORKDIR}/{${pkgdir}/{debian,NVIDIA-Linux,NVIDIA-Linux-amd64,NVIDIA-Linux-x86_64-${pkgver},LICENSE.txt,unpack-stamp},compiled_deb,compiled_other}

###########################################################

COMMANDS=(
  apt
  dpkg
  grep
  sudo
  wc
  wget
)

function checkCommands() {

    if [[ $(which --help 2>/dev/null) ]] && [[ $(echo --help 2>/dev/null) ]]; then

        local a=0
        for command in ${@}; do
            if [[ ! $(which $command 2>/dev/null) ]]; then
                local COMMANDS_NOTFOUND[$a]=${command}
                let a++
            fi
        done

        if [[ -n $COMMANDS_NOTFOUND ]]; then
            echo -e "\nError! The following commands could not be found: ${COMMANDS_NOTFOUND[*]}\nAborting\n"
            exit 1
        fi
    else
        exit 1
    fi
}

checkCommands "${COMMANDS[*]}"

########################################################

# General function for question responses

function questionresponse() {

  local response=${1}

  read -r -p "" response
  if [[ $(echo $response | sed 's/ //g') =~ ^([yY][eE][sS]|[yY])$ ]]; then
    echo ""
    return 0
  else
    return 1
  fi

}

########################################################

# Premilinary check for already compiled packages

if [[ $(ls ${WORKDIR}/*.deb 2>/dev/null | wc -l) -ne 0 ]]; then
  echo -e "\nWarning: previously compiled deb archives found on the main build directory '${WORKDIR}'.\nDelete them and continue? [Y/n]"
  questionresponse
fi

if [[ $? -ne 0 ]]; then
  echo -e "Cancelling.\n"
  exit 1
else
  rm -rv ${WORKDIR}/*.{deb,buildinfo,changes} 2>/dev/null
fi

########################################################

# Premilinary check to see whether Nvidia card is present

if [[ ! $(lspci | grep -oiE "vga.*nvidia") ]]; then
  echo -e "\nWarning: Nvidia card could not be detected on your system. Continue anyway? [Y/n]"
  questionresponse
fi

if [[ $? -ne 0 ]]; then
  echo -e "Cancelling.\n"
  exit 1
fi

########################################################

# Premilinary check to see whether Nvidia card is present

echo -e "\nAuto-install Nvidia drivers after compilation? [Y/n]"
questionresponse

if [[ $? -eq 0 ]]; then

  if [[ ${UID} -ne 0 ]]; then
    if [[ $(echo $(sudo -vn &>/dev/null)$?) -ne 0 ]]; then
      echo -e "Installation requires root permissions. Please provide your sudo password now.\n"
      sudo -v
    fi
    if [[ $? -eq 0 ]]; then
      AUTOINSTALL=
    else
      exit 1
    fi
  else
    AUTOINSTALL=
  fi
fi

########################################################

# Check and install package related dependencies if they are missing
function pkgdependencies() {

  # Generate a list of missing dependencies
  local a=0
  for p in ${@}; do
    if [[ $(echo $(dpkg -s ${p} &>/dev/null)$?) -ne 0 ]]; then
      validlist[$a]=${p}
      let a++
    fi
  done

  if [[ -n ${list[*]} ]]; then
    echo -e "Some build time dependencies are missing. In order to continue, you must install them."
    if [[ ${UID} -ne 0 ]] && [[ $(echo $(sudo -vn &>/dev/null)$?) -ne 0 ]]; then
      echo -e "For that, sudo password is required. Please provide it now.\n"
      sudo -v
      if [[ $? -ne 0 ]]; then
        echo -e "Error: couldn't continue due to lacking permissions. Aborting\n"
        exit 1
      fi
    fi
  fi

  # Install missing dependencies, be informative
  local b=0
  for pkgdep in ${validlist[@]}; do
    echo -e "$(( $b + 1 ))/$(( ${#validlist[*]} )) - Installing ${_pkgname} dependency ${pkgdep}"
    sudo apt install -y ${pkgdep} &> /dev/null
    if [[ $? -eq 0 ]]; then
      let b++
    else
      echo -e "\nError occured while installing ${pkgdep}. Aborting.\n"
      exit 1
    fi
  done
  if [[ -n ${validlist[*]} ]]; then
    # Add empty newline
    echo ""
  fi
}

########################################################

function download_files() {

  mkdir -p ${WORKDIR}/${pkgdir}

  i=0
  for f in ${files[@]}; do

    if [[ ! -f ${WORKDIR}/${pkgdir}/${filebases[$i]} ]]; then

      echo -e "\nDownloading ${filebases[$i]}"
      wget ${f} -o /dev/null --show-progress -O "${WORKDIR}/${pkgdir}/${filebases[$i]}"

      if [[ $? -ne 0 ]];then
        echo -e "Error: couldn't retrieve file ${filebases[$i]}. Aborting\n"
        exit 1
      fi
    fi

  let i++
  done

}

########################################################

function prepare_deb_sources() {

  # Extract debian control files
  cd ${WORKDIR}/${pkgdir}/
  tar xf ${filebases[0]}

  if [[ $? -eq 0 ]]; then

    function fix_library_versions() {

      for oldlib in ${!library_fixes[@]}; do

        # sed-friendly name
        local oldlib_sed=$(printf '%s' ${oldlib} | sed 's/\./\\\./g')

        for lib in ${library_fixes[$oldlib]}; do

          # sed-friendly name
          local lib_sed=$(printf '%s' ${lib} | sed 's/\./\\\./g')

          # Files which have old library files mentioned
          local i=0
          for oldlib_file in $(grep -rl "${oldlib}" debian/ | tr '\n' ' '); do
            local oldlib_files[$i]=${oldlib_file}
            let i++
          done

          for targetfile in ${oldlib_files[@]}; do
            sed -i "s/${oldlib_sed}/${lib_sed}/g" ${targetfile}
          done
        done
      done

    }

    function rename_deb_files() {

      # Remove this suffix
      sed -i 's/\-no\-compat32//' debian/rules.defs

      # Tell that Nvidia .run file is at our build root, not in amd64 subfolder
      sed -i 's|sh \$\*\/\${NVIDIA_FILENAME_\$\*}|sh \${NVIDIA_FILENAME_\$\*}|' debian/rules

      ############
      # TODO Individual fix for strange version number present in debian control files
      # Remove when not needed!
      sed -i "s/384/${pkgver_major}/g" debian/control
      sed -i "s/384/${pkgver_major}/g" debian/templates/control.in
      ############

      local IFS=$'\n'
      for n in $(ls debian/ -w 1); do

        # IMPORTANT! KEEP THIS IF STATEMENT ORDER BELOW!!

        # Do this for every file in debian subfolder regardless of their name
        if [[ -f debian/${n} ]]; then
          # Keep this order. It is important!
          sed -i "s/${oldver_sed}/${pkgver_sed}/g" debian/${n}
          sed -i "s/${oldver_major}/${pkgver_major}/g" debian/${n}

        fi

        if [[ $(printf '%s' ${n} | grep ${oldver_major}) ]]; then
          local n_new=$(printf '%s' ${n} | sed "s/${oldver_major}/${pkgver_major}/")
          mv debian/${n} debian/${n_new}

          if [[ $? -ne 0 ]]; then
            echo -e "Error: couldn't rename file debian/${n}. Aborting\n"
            exit 1
          fi
        fi

      done
      unset IFS
    }

    fix_library_versions
    rename_deb_files

  else
    echo -e "Error: couldn't extract Nvidia Debian archive. Aborting\n"
    exit 1
  fi

}

########################################################

function compile_nvidia() {

  cd ${WORKDIR}/${pkgdir}/
  DEB_BUILD_OPTIONS="strip noddebs" dpkg-buildpackage -rfakeroot -b -us -uc

  if [[ $? -eq 0 ]]; then
    mkdir -p ${WORKDIR}/compiled_deb
    for p in ${nvidia_install_packages[*]}; do
      mv ${WORKDIR}/${p}*.deb ${WORKDIR}/compiled_deb/
    done

    if [[ $? -eq 0 ]]; then
      echo -e "Compiled deb packages moved into '${WORKDIR}/compiled_deb/'. Install these to enable Nvidia support.\nAdditionally, you may need Vulkan loader package 'libvulkan1', too.\n"
    else
      echo -e "Error: couldn't move deb packages into '${WORKDIR}/compiled_deb/'. Aborting\n"
      exit 1
    fi

    mkdir -p ${WORKDIR}/compiled_other
    for a in ${WORKDIR}/*; do
      if [[ -f ${a} ]]; then
        mv ${a} ${WORKDIR}/compiled_other/
      fi
    done

    if [[ $? -eq 0 ]]; then
      echo -e "Other files moved into '${WORKDIR}/compiled_other/' \n"
    else
      echo -e "Error: couldn't move other files into '${WORKDIR}/compiled_other/'. Aborting\n"
      exit 1
    fi

  else
    echo -e "Error: couldn't compile Nvidia package bundle. Aborting\n"
    exit 1
  fi
}

########################################################

function install_nvidia() {

  for syspkg in ${nvidia_required_packages[@]}; do
    if [[ $(echo $(dpkg -s ${syspkg} &>/dev/null)$?) -ne 0 ]]; then
      echo -e "Installing missing dependency ${syspkg}\n"
      sudo apt install -y ${syspkg}
      if [[ $? -ne 0 ]]; then
        echo -e "Error: couldn't install dependency ${syspkg}. Aborting\n"
        exit 1
      fi
    fi
  done

  cd ${WORKDIR}/compiled_deb
  for pkg in ${nvidia_install_packages[@]}; do

    local oldpkg=$(printf '%s' ${pkg} | sed 's/\-[0-9]*$//')
    local oldpkg_check=$(dpkg --get-selections | grep ${oldpkg} | awk '{print $1}')

    if [[ $(echo ${oldpkg_check} | wc -w) -eq 1 ]]; then
      if [[ ! ${oldpkg_check} =~ ^*${pkgver_major}$ ]]; then
        echo -e "Removing old ${oldpkg}\n"
        sudo apt purge --remove -y "${oldpkg}*"
        if [[ $? -ne 0 ]]; then
          echo -e "Error: couldn't uninstall ${oldpkg}. Aborting\n"
          exit 1
        fi
      fi
    fi

    echo -e "Installing ${pkg}\n"
    sudo dpkg -i ${pkg}*.deb

    if [[ $? -ne 0 ]]; then
      echo -e "Warning: couldn't install ${pkg}\n"
    fi

  done

}

function install_vulkan() {

  # Vulkan loader
  if [[ $? -eq 0 ]]; then
    local syspkg=libvulkan1
    if [[ $(echo $(dpkg -s ${syspkg} &>/dev/null)$?) -ne 0 ]]; then
      sudo apt update && sudo apt install -y ${syspkg}
    fi
  fi

}

########################################################

download_files && \
pkgdependencies ${nvidia_builddeps[*]}
prepare_deb_sources && \
compile_nvidia

if [[ $? -eq 0 ]] && [[ -v AUTOINSTALL ]]; then
  install_nvidia && \
  install_vulkan
fi

# If any buildtime deps were installed, inform the user
if [[ -n ${list[*]} ]]; then
  echo -e "The following buildtime dependencies were installed, and they may not be required anymore:\n\n\
$(for h in ${list[*]}; do echo ${h}; done)\n"
fi