#!/usr/bin/bash

CMDLINE_ARGS_NVIDIA="rd.driver.blacklist=nouveau"
CMDLINE_ARGS_ALWAYS_REMOVE="nomodeset gfxpayload=vga=normal nouveau.modeset=0 nvidia-drm.modeset=1 nvidia-drm.fbdev=1 modprobe.blacklist=nouveau initcall_blacklist=simpledrm_platform_driver_init"

print_usage() {
cat <<EOF
Tool to add or remove kernel command line options required for proper operation of the Nvidia driver.
Its main use is to be called from the %post/%preun scripts of the Nvidia driver packages, but it can also be used in other contexts, for example in a kickstart file after the drivers have been already installed.

Usage: nvidia-update-boot post|preun
    post        Adjust necessary kernel command line options
    preun       Remove all kernel command line options

EOF
}

check() {
  if [ -f /etc/default/grub ]; then

    # Grub 2 is installed
    if [ -d /sys/firmware/efi ]; then
      GRUB_CFG=/etc/grub2-efi.cfg
    else
      GRUB_CFG=/etc/grub2.cfg
    fi

    . /etc/default/grub

  else

    echo "Nvidia driver setup: no bootloader configured. Please run 'nvidia-boot-update post' manually."

  fi
}

post() {

  # Edit GRUB configuration file
  if [ -v GRUB_CFG ]; then

    if [ -z "${GRUB_CMDLINE_LINUX}" ]; then
      echo GRUB_CMDLINE_LINUX=\""$CMDLINE_ARGS_NVIDIA"\" >> /etc/default/grub
    else
      for param in $CMDLINE_ARGS_NVIDIA; do
        echo ${GRUB_CMDLINE_LINUX} | grep -q $param
        [ $? -eq 1 ] && GRUB_CMDLINE_LINUX="${GRUB_CMDLINE_LINUX} ${param}"
      done
      for param in $CMDLINE_ARGS_ALWAYS_REMOVE; do
        echo ${GRUB_CMDLINE_LINUX} | grep -q $param
        [ $? -eq 0 ] && GRUB_CMDLINE_LINUX="$(echo ${GRUB_CMDLINE_LINUX} | sed -e "s/ $param//g")"
      done
      sed -i -e "s|^GRUB_CMDLINE_LINUX=.*|GRUB_CMDLINE_LINUX=\"${GRUB_CMDLINE_LINUX}\"|g" /etc/default/grub
    fi

    grub2-mkconfig -o $GRUB_CFG &>/dev/null

  fi

}

preun() {

  # Edit GRUB configuration file
  if [ -v GRUB_CFG ]; then

    for param in $CMDLINE_ARGS_NVIDIA; do
      echo ${GRUB_CMDLINE_LINUX} | grep -q $param
      [ $? -eq 0 ] && GRUB_CMDLINE_LINUX="$(echo ${GRUB_CMDLINE_LINUX} | sed -e "s/ $param//g")"
    done
    sed -i -e "s|^GRUB_CMDLINE_LINUX=.*|GRUB_CMDLINE_LINUX=\"${GRUB_CMDLINE_LINUX}\"|g" /etc/default/grub

    grub2-mkconfig -o $GRUB_CFG &>/dev/null

  fi

}

case "$1" in
  post)
    check
    post
    exit 0
    ;;
  preun)
    check
    preun
    exit 0
    ;;
  *)
    print_usage
    exit 0
    ;;
esac
