#!/bin/bash
set -e
shopt -s nullglob

NV_CURRENT=/usr/lib/nvidia/current
NV_KERNEL_OBJ=${NV_CURRENT}/kernel-obj
MODULE_DIR=/var/lib/endless-external-drivers/nvidia
DATA_DIR=/usr/share/endless-external-drivers/nvidia
DATA_OVERRIDE_DIR=/etc/endless-external-drivers/nvidia

NV_DEVICE=$(lspci -mn | awk '{ gsub("\"",""); if (($2 == "0300" || $2 == "0302") && ($3 == "10de" || $3 == "12d2")) { print toupper($3)toupper($4) } }' | head -n 1)
NV_SUBVENDOR=$(lspci -mn | awk '{ gsub("\"",""); if (($2 == "0300" || $2 == "0302") && ($3 == "10de" || $3 == "12d2")) { print toupper($6) } }' | head -n 1)

if [[ -z ${NV_DEVICE} ]]; then
  exit 0
fi

echo "Found nvidia device ${NV_DEVICE}"

# Return 1 if nvidia is blacklisted
nvidia_blacklist_check_by_product() {
  # Try a user-defined dmi blacklist override, which would not ordinarily be
  # present
  local dmi_blacklist=${DATA_OVERRIDE_DIR}/dmi-blacklist
  if ! [[ -e ${dmi_blacklist} ]]; then
    # Fall back on the ostree-shipped DMI blacklist
    dmi_blacklist=${DATA_DIR}/dmi-blacklist
  fi

  # Accurately parsing CSV in bash is impractical. Use python to do this right.
  if [[ -e ${dmi_blacklist} ]]; then
    python3 -c "
import csv, sys
sys_vendor = open('/sys/class/dmi/id/sys_vendor').read().strip()
product_name = open('/sys/class/dmi/id/product_name').read().strip()
with open('${dmi_blacklist}', newline='') as fp:
    # Filter out empty lines and comments
    reader = csv.reader(row for row in fp if row.rstrip()
                        and not row.startswith('#'))
    sys.exit([sys_vendor, product_name] in reader)
" || return 1
  fi

  return 0
}

# Return 1 if nvidia is blacklisted
nvidia_blacklist_check_by_board() {
  # Try a user-defined dmi blacklist override, which would not ordinarily be
  # present
  local dmi_board_blacklist=${DATA_OVERRIDE_DIR}/dmi-board-blacklist
  if ! [[ -e ${dmi_board_blacklist} ]]; then
    # Fall back on the ostree-shipped DMI blacklist
    dmi_board_blacklist=${DATA_DIR}/dmi-board-blacklist
  fi

  # Accurately parsing CSV in bash is impractical. Use python to do this right.
  if [[ -e ${dmi_board_blacklist} ]]; then
    python3 -c "
import csv, sys
board_vendor = open('/sys/class/dmi/id/board_vendor').read().strip()
board_name = open('/sys/class/dmi/id/board_name').read().strip()
with open('${dmi_board_blacklist}', newline='') as fp:
    # Filter out empty lines and comments
    reader = csv.reader(row for row in fp if row.rstrip()
                        and not row.startswith('#'))
    sys.exit([board_vendor, board_name] in reader)
" || return 1
  fi

  return 0
}

# Return 1 if nvidia is blacklisted
nvidia_blacklist_check_by_pci_id() {
  # Try a blacklist override specific to ASUS PCI IDs
  local asus_pci_blacklist=${DATA_OVERRIDE_DIR}/asus-pci-blacklist
  if ! [[ -e ${asus_pci_blacklist} ]]; then
    # Fall back on the ostree-shipped ASUS PCI blacklist
    asus_pci_blacklist=${DATA_DIR}/asus-pci-blacklist
  fi

  if [[ -e ${asus_pci_blacklist} ]] &&
     [[ "1043" == "${NV_SUBVENDOR}" ]]; then
    grep --quiet -F ${NV_DEVICE} ${asus_pci_blacklist} && return 1
  fi

  return 0
}

# Return 1 if module name passed as argument is blacklisted through the kernel
# command line with modprobe.blacklist. Note that modprobe.blacklist accepts a
# comma-separated list of modules, i.e. modprobe.blacklist=nouveau,nvidia.
check_cmdline_blacklist() {
  module=${1}
  set -- $(< /proc/cmdline)
  for param in "$@"; do
      case "$param" in
          modprobe.blacklist=*)
          if [[ ",${param#modprobe.blacklist=}," =~ ",${module}," ]]; then
              echo "${module} blacklisted on kernel command line"
              return 1
          fi
          ;;
      esac
  done

  return 0
}

# Return 1 if nvidia is blacklisted
nvidia_blacklist_check() {
  if ! check_cmdline_blacklist "nvidia" ||
     ! nvidia_blacklist_check_by_product ||
     ! nvidia_blacklist_check_by_board ||
     ! nvidia_blacklist_check_by_pci_id; then
    return 1
  fi

  return 0
}

# Load a module from a .ko file after first trying to load it's dependencies
# using modprobe.
load_module() {
	local deps=$(modinfo -F depends "$1")
	if [[ -n $deps ]]; then
		modprobe -a -q ${deps//,/ } || :
	fi
	insmod "$@" || :
}

# Build the nvidia module if we haven't already built the current
# ostree-shipped driver version for the currently running kernel.
build_nvidia_if_needed() {
  local running_kernel_version=$(uname -r)
  local nvidia_ostree_version
  local built_driver_version
  local built_kernel_version

  if [[ -e "${NV_KERNEL_OBJ}/version" ]]; then
    nvidia_ostree_version=$(<${NV_KERNEL_OBJ}/version)
  fi
  if [[ -e "${MODULE_DIR}/driver-version" ]]; then
    built_driver_version=$(<${MODULE_DIR}/driver-version)
  fi
  if [[ -e "${MODULE_DIR}/kernel-version" ]]; then
    built_kernel_version=$(<${MODULE_DIR}/kernel-version)
  fi

  echo "Kernel version: running=${running_kernel_version} built=${built_kernel_version}"
  echo "Driver version: shipped=${nvidia_ostree_version} built=${built_driver_version}"

  # If we've built the driver for the kernel we're running, and it's the same
  # driver version as the one in the ostree, then we have nothing to do
  if [[ "${built_kernel_version}" == "$(uname -r)" ]] &&
     [[ "${built_driver_version}" == "${nvidia_ostree_version}" ]]; then
    return
  fi

  # Otherwise, build needed.
  echo "Building driver"
  rm -rf "${MODULE_DIR}"
  mkdir -p "${MODULE_DIR}"
  OUTPREFIX="${MODULE_DIR}/" "${NV_KERNEL_OBJ}/build" || return

  # Append signatures, if we have them
  local kmod
  for kmod in ${MODULE_DIR}/*.ko; do
    cat "${NV_KERNEL_OBJ}/$(basename $kmod).sig" >> ${kmod} 2>/dev/null || :
  done

  uname -r > "${MODULE_DIR}/kernel-version"
  cp "${NV_KERNEL_OBJ}/version" "${MODULE_DIR}/driver-version"
}

# If no user accounts are available, assume we're still in the factory
# and hence properietary drivers should be erased
first_boot_cleanup_modules()
{
  local homedirs=( /home/* )
  if [[ ${#homedirs[@]} == 0 ]]; then
    rm -rf /var/lib/endless-external-drivers/nvidia
  fi
}

if ! modprobe -c | grep -F -x --quiet "blacklist nvidia" &&
   grep --quiet -F ${NV_DEVICE} ${NV_CURRENT}/nvidia.ids &&
   nvidia_blacklist_check; then

  build_nvidia_if_needed
  echo "Loading nvidia modules"
  load_module "${MODULE_DIR}"/nvidia.ko NVreg_PreserveVideoMemoryAllocations=1
  load_module "${MODULE_DIR}"/nvidia-modeset.ko
  load_module "${MODULE_DIR}"/nvidia-drm.ko modeset=1
  first_boot_cleanup_modules
  if [[ -e "/sys/module/nvidia_drm" ]]; then
    echo "nvidia loaded successfully"
    exit 0
  fi

  # Try to unload anything that loaded
  echo "nvidia load failed"
  rmmod nvidia-drm || :
  rmmod nvidia-modeset || :
  rmmod nvidia || :
fi

# nouveau selected, or we want to use nouveau because nvidia build/load failed

# nouveau claims all devices based on the nvidia vendor IDs so we don't
# need to do a specific product ID lookup here, just load it.
if check_cmdline_blacklist "nouveau"; then
  echo "Loading nouveau"
  modprobe nouveau
fi
