#!/usr/bin/env bash
set -o errexit
set -o nounset
set -o pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
# shellcheck source=00-platform-guard.inc
source "$SCRIPT_DIR/00-platform-guard.inc"

driver_version=""

usage() {
  cat <<'EOF'
Usage: sudo ./06-nvidia-tools.sh [--install-driver VERSION]

Without --install-driver, only non-driver diagnostic tools are installed.
EOF
}

while (($# > 0)); do
  case "$1" in
    --install-driver)
      if (($# < 2)); then
        printf 'CRITICAL: --install-driver requires a VERSION\n' >&2
        exit 2
      fi
      driver_version="$2"
      if [[ ! "$driver_version" =~ ^[0-9]+$ ]]; then
        printf 'CRITICAL: NVIDIA driver VERSION must contain digits only\n' >&2
        exit 2
      fi
      shift
      ;;
    -h|--help)
      usage
      exit 0
      ;;
    *)
      printf 'CRITICAL: unknown option: %s\n' "$1" >&2
      exit 2
      ;;
  esac
  shift
done

if ((EUID != 0)); then
  printf 'CRITICAL: NVIDIA tooling setup must run as root\n' >&2
  exit 2
fi
require_supported_ubuntu
if ! command -v apt-get >/dev/null 2>&1; then
  printf 'CRITICAL: apt-get is required\n' >&2
  exit 2
fi

apt-get update
DEBIAN_FRONTEND=noninteractive apt-get install -y nvtop clinfo pciutils

printf '\n== NVIDIA PCI devices ==\n'
lspci -nn | grep -i nvidia || printf 'INFO: no NVIDIA PCI devices detected\n'

printf '\n== NVIDIA runtime ==\n'
if command -v nvidia-smi >/dev/null 2>&1; then
  nvidia-smi || printf 'WARNING: nvidia-smi returned an error\n'
else
  printf 'INFO: nvidia-smi is not installed\n'
fi

printf '\n== DKMS ==\n'
if command -v dkms >/dev/null 2>&1; then
  dkms status || printf 'WARNING: dkms status returned an error\n'
else
  printf 'INFO: dkms is not installed\n'
fi

if [[ -n "$driver_version" ]]; then
  driver_package="nvidia-driver-$driver_version"
  if ! apt-cache show "$driver_package" >/dev/null 2>&1; then
    printf 'CRITICAL: requested NVIDIA driver package is unavailable: %s\n' \
      "$driver_package" >&2
    exit 1
  fi
  DEBIAN_FRONTEND=noninteractive apt-get install -y "$driver_package"
  printf 'WARNING: NVIDIA driver %s was installed; reboot before validation\n' \
    "$driver_version"
else
  printf 'OK: NVIDIA diagnostic tools installed; no driver was installed\n'
fi