#!/bin/bash

# On Nvidia GPU machines, the non-dkms Nvidia drivers need to be updated
# *in sync* with the corresponding kernel:
#   - if linux-lts is updated and nvidia-lts is installed, nvidia-lts should be updated too
#   - if linux     is updated and nvidia     is installed, nvidia     should be updated too
#
# This app checks the above and informs the calling app with exit code:
#     0 = success, no issue detected
#     1 = failure, update is missing
#     2 = usage error
# On failure also the appopriate error messages will be displayed to stderr.

DIE() { echo "==> $progname: error: $1" >&2; exit 2; }

ExitOK()   { [ $verbose = yes ] && echo "==> no Nvidia driver update issue detected. " >&2; exit 0; }
ExitFail() { exit 1; }

UpdatesInclude() { echo "$updates" | grep "^$1$" >/dev/null ; }
IsInstalled()    { expac -Q %n "$1" >/dev/null ; }
DriverCheck()    {
    local -r linux="$1"
    local -r nvidia="$2"
    if UpdatesInclude "$linux" && IsInstalled "$nvidia" && ! UpdatesInclude "$nvidia" ; then
        msgs+=("updates include $linux but not $nvidia")
    fi
}

Parameters() {
    while true ; do
        case "$1" in
            -v | --verbose) verbose=yes ;;
            --no-color)     RED=""; RESET="" ;;
            -*)             DIE "unsupported option '$1'" ;;
            *)              if [ "$1" ] ; then
                                updates="$(printf "%s\n" "$@")"                  # from parameters
                            else
                                updates="$(checkupdates | awk '{print $1}')"     # from a command
                            fi
                            return
                            ;;
        esac
        shift
    done
}

Main() {
    expac %n nvidia nvidia-lts >/dev/null || exit 0     # nothing to do

    local -r progname=${0##*/}
    local RED=$'\e[0;91m'
    local RESET=$'\e[0m'
    local verbose=no
    local updates=""             # will contain all updates, separated by newline; see Parameters()
    local msgs=()

    echo ":: Nvidia check..." >&2

    Parameters "$@"

    DriverCheck linux     nvidia
    DriverCheck linux-lts nvidia-lts

    if [ ${#msgs[@]} -eq 0 ] ; then
        ExitOK
    else
        echo "${RED}==> $progname: warning: $(printf "%s\n" "${msgs[@]}")${RESET}" >&2
        ExitFail
    fi
}

Main "$@"
