#!/bin/bash
#
# Updates NVIDIA database for nvidia-installer and nvidia-installer-dkms.
#
# https://www.nvidia.com/en-us/drivers/unix/
# https://download.nvidia.com/XFree86/Linux-x86_64/435.21/README/supportedchips.html
# https://download.nvidia.com/XFree86/Linux-x86_64/435.21/README/index.html


source /usr/share/endeavouros/scripts/eos-script-lib-yad || {
    echo "Error: $(basename "$0"): file eos-script-lib-yad not found." >&2
    exit 1
}

LATEST_BRANCH1="Latest Production Branch Version"         # "Latest Long Lived Branch"
LATEST_BRANCH2="Latest New Feature Branch Version"        # "Latest Short Lived Branch"

LATEST_LEGACY="Latest Legacy GPU version"

GetNvidiaBranch() {
    local string="$1"
    local series="$2"
    local result="$(curl -Lsm 10 -o- https://www.nvidia.com/en-us/drivers/unix | sed 's|</a>|</a>\n|g' | grep "$string" | head -n1 | sed 's|.*>\([0-9.]*\)[ ]*<.*|\1|')"

    test -n "$result" || DIE "$FUNCNAME: fetching branch '$string' failed."
    Debug "$series branch: $result"
    echo "$result"
}

GetNvidiaIds() {
    local branch="$1"
    local ids="" #"$(curl -Lsm 10 -o- https://download.nvidia.com/XFree86/Linux-x86_64/$branch/README/supportedchips.html)"
    [ -n "$ids" ] || ids="$(curl -Lsm 10 -o- https://us.download.nvidia.com/XFree86/Linux-x86_64/$branch/README/supportedchips.html)"
    test -n "$ids" || DIE "$FUNCNAME: fetching ids from branch '$branch' failed."
    
    ids="$(echo "$ids" | sed -n '1,/.*Below are the legacy GPUs that are no longer supported.*/p' | grep "</td>" | grep "<td>")"
    ids="$(echo "$ids" | grep "<td>[0-9A-F][0-9A-F][0-9>-F][0-9A-F][ <]" | sed 's|<td>\(....\).*|\1|' | sort | uniq | tr '[:upper:]' '[:lower:]')"
    echo "$ids"
}

NvidiaIdsCommon() {
    local string="$1"
    local series="$2"
    local branch="$3"
    [ -n "$branch" ] || branch="$(GetNvidiaBranch "$string" "$series")"
    test -n "$branch" || DIE "$FUNCNAME: GetNvidiaBranch failed."
    if [ "$series" = "latest" ] && [ "$branch" != "$pacman_latest_branch" ] ; then
        cat <<EOF >&2

Warning from $PROGNAME:
    The nvidia.com site has database of card ids for the new driver ($branch).
    Thus the nvidia.com database will reflect the new driver.
    The driver from pacman has a different version ($pacman_latest_branch).
    It is recommended to run $PROGNAME again after having
    the pacman version updated to the nvidia.com version.
EOF
        if [ "$targetdir" != "/tmp" ] ; then
            read -p "Want to proceed with the db update (Y/n)? "
            case "$REPLY" in
                ""|[yY]*) ;;
                *) return 1 ;;
            esac
        fi
    fi
    local ids="$(GetNvidiaIds "$branch")"
    test -n "$ids" || DIE "$FUNCNAME: GetNvidiaIds failed."
    Info "$string: $branch"
    echo "$ids"
}

LatestNvidiaBranch() {
    local br br1="" br2=""

    br1="$(GetNvidiaBranch "$LATEST_BRANCH1" "latest")"
    if [ "$br1" != "$pacman_latest_branch" ] ; then
        br2="$(GetNvidiaBranch "$LATEST_BRANCH2" "latest")"
    fi
    case "$pacman_latest_branch" in
        "$br1") br="$br1"
                latest_branch_string_used="$LATEST_BRANCH1"
                ;;
        "$br2") br="$br2"
                latest_branch_string_used="$LATEST_BRANCH2"
                ;;
        *)
            # use the latest of the available databases
            local vercmp=$(vercmp "$br1" "$br2")
            if [ $vercmp -ge 0 ] ; then
                br="$br1"
                latest_branch_string_used="$LATEST_BRANCH1"
            else
                br="$br2"
                latest_branch_string_used="$LATEST_BRANCH2"
            fi
            ;;
    esac

    latest_available_nvidia_branch="$br"
}

NvidiaIdsLatest() {
    local ids

    ids="$(NvidiaIdsCommon "$latest_branch_string_used" "latest" "$latest_available_nvidia_branch")"
    if [ -n "$ids" ] ; then
        Info "$latest_branch_string_used: found!"
        echo "$ids"
        return 0
    else
        return 1
    fi
    DIE "$FUNCNAME: NvidiaIdsCommon failed."
}

NvidiaIdsLegacy() {
    local series="$legacy_series.xx series"
    local ids="$(NvidiaIdsCommon "$LATEST_LEGACY" "$series")"
    test -n "$ids" || DIE "$FUNCNAME: NvidiaIdsCommon failed for Legacy $series."
    Info "$LATEST_LEGACY: found!"
    echo "$ids"
}

echo2()   { echo "$@" >&2 ; }
DIE()     { echo2 "$PROGNAME: error: $1" ; exit 1 ; }
WARN()    { echo2 "$PROGNAME: warning: $1" ; }
Info()    { test "$show_additional_info" = "yes" && echo2 "====> $@" ; }
Debug()   { test "$show_additional_info_debug" = "yes" && echo2 "====> $@" ; }

Options() {
    local arg
    for arg in "$@" ; do
        case "$arg" in
            -h | --help)
                cat <<EOF >&2
Usage: $PROGNAME [options]
Options:
    --help    | -h  This help.
    --verbose | -v  Shows some additional info messages.
    --debug)        Shows some additional info and debug messages.
    --tmpdb)        Saves card id file(s) into the /tmp folder.
    --here)         Saves card id file(s) into current folder.
EOF
                exit 0
                ;;
            -v | --verbose)
                show_additional_info=yes ;;
            --debug)
                show_additional_info=yes
                show_additional_info_debug=yes ;;
            --tmpdb)
                targetdir="/tmp" ;;
            --here)
                targetdir="." ;;
            -*)
                DIE "unsupported option '$arg'" ;;
            *)
                DIE "unsupported parameter '$arg'" ;;
        esac
    done
}

ExeEnv() {
    case "$(echo "${BASH_SOURCE[@]}")" in
        *rogalmic.bash-debug*)  exe_env=vscode ;;
        */usr/bin/bashdb*)      exe_env=bashdb ;;
        *)                      exe_env=run ;;
    esac
}

Logging() {
    local newfile="$1"
    local currfile="$2"
    local branch="$3"
    local xx

    if [ ! -r "$currfile" ] ; then
        printf "\nBranch $branch: new db created!\n" >> $log
        has_logs=yes
    else
        xx="$(diff $newfile $currfile)"
        if [ -n "$xx" ] ; then
            printf "\nDiffs for the $branch branch:\n" >> $log
            echo "$xx" >> $log
            has_logs=yes
        fi
    fi
}

CreateNvidiaDb()
{
    local PROGNAME=nvidia-installer-update-db
    local exe_env=""
    local show_additional_info=no                       # "no" is the default, "yes" is TESTING!
    local show_additional_info_debug=no                 # "no" is the default, "yes" is TESTING!
    local targetdir=/var/lib/pci                        # db directory

    Options "$@"

    local legacy_series="390"                           # this may change some day!
    local legacy_available=yes                          # -"-

    local target_latest=$targetdir/nvidia.ids           # db (list of ids) for driver type 'nvidia'
    local target_legacy=$targetdir/nvidia-${legacy_series}xx.ids     # db (list of ids) for legacy driver type 'nvidia-NNNxx'
    local bakdir=$targetdir/ORIG.BAK                    # original db is stored into this directory, if it exists
    local nvidia_ids_latest
    local nvidia_ids_legacy
    local tmpfile1
    local tmpfile2
    local sudo_commands=":"                             # collect required root commands here
    #local readme="NVIDIA db generated by $PROGNAME at $(date +%Y-%m-%d)."
    local xx
    local log=$(mktemp /tmp/$PROGNAME.log.XXXXX)
    local has_logs=no
    local pacman_latest_branch=$(LANG=C pacman -Si nvidia-dkms | grep ^Version | head -n1 | awk '{print $NF}' | sed 's|-[0-9][0.9]*$||')
    local legacy_version=""
    local latest_branch_string_used=""
    local latest_available_nvidia_branch=""

    LatestNvidiaBranch

    Debug "pacman version: $pacman_latest_branch"

    if [ "$targetdir" = "/tmp" ] ; then
        rm -f $target_latest $target_legacy  # option --tmpdb: always create new db into /tmp
    fi

    # check that legacy driver is available either in the repos or AUR
    if [ -z "$(yay -S -s nvidia -s dkms -s "$legacy_series"xx -q)" ] ; then
        legacy_available=no  # not available
    else
        legacy_version=$(yay -Si nvidia-${legacy_series}xx-dkms | grep ^Version | awk '{print $3}' | sed 's|-[0-9][0-9]*$||')
        Debug "legacy version: $legacy_version"
    fi

    Info "Fetching NVIDIA graphics card ids from nvidia.com ..."

    nvidia_ids_latest="$(NvidiaIdsLatest)"
    [ $? -eq 0 ] || return 1
    test -n "$nvidia_ids_latest" || DIE "cannot fetch latest NVIDIA ids"

    if [ "$legacy_available" = "yes" ] ; then
        nvidia_ids_legacy="$(NvidiaIdsLegacy)"
        test -n "$nvidia_ids_legacy" || DIE "cannot fetch legacy NVIDIA Legacy $legacy_series.xx ids"
    fi

    # Data is fetched from nvidia.com.
    # Check if anything has changed.

    tmpfile1=$(mktemp)
    echo "$nvidia_ids_latest" > $tmpfile1
    Logging "$tmpfile1" "$target_latest" "Latest"
    
    if [ "$legacy_available" = "yes" ] ; then
        tmpfile2=$(mktemp)
        echo "$nvidia_ids_legacy" > $tmpfile2
        Logging "$tmpfile2" "$target_legacy" "Legacy"
    fi

    if [ "$has_logs" = "no" ] ; then
        rm -f $log
        echo "Database is already up-to-date." >&2
        if [ "$targetdir" != "/tmp" ] ; then
            test -r $tmpfile1 && rm -f $tmpfile1
            test -r $tmpfile2 && rm -f $tmpfile2
            return
        fi
    fi

    local date="Date & time: $(date -u "+%Y-%m-%d %H:%M") UTC  [yyyy-mm-dd hh:mm, 24h clock]"
    local metainfo="$date\nNVIDIA version: $latest_available_nvidia_branch\nPacman version: $pacman_latest_branch\n"

    # Now save the data into the local database.

    case "$targetdir" in
        .)
            # save to current working folder
            for xx in $target_latest $target_legacy ; do
                if [ -r $xx ] ; then
                    rm -f $xx
                else
                    WARN "$xx does not exist."
                fi
            done
            mv $tmpfile1 $target_latest
            mv $tmpfile2 $target_legacy
            printf "$metainfo" > $target_latest.metainfo
            chmod 0644 $target_latest $target_legacy $target_latest.metainfo
            return
            ;;
    esac

    # $targetdir is /var/lib/pci or /tmp

    sudo_commands+=" ; mkdir -p $targetdir"
    sudo_commands+=" ; if [ ! -d $bakdir ] ; then mkdir -p $bakdir ; cp -a $targetdir/*.ids $bakdir >& /dev/null ; fi"
    sudo_commands+=" ; mv $tmpfile1 $target_latest"
    if [ "$targetdir" != "/tmp" ] ; then
        sudo_commands+=" ; chown root:root $target_latest"
    fi
    sudo_commands+=" ; printf '$metainfo' > $target_latest.metainfo"
    sudo_commands+=" ; chmod 0644 $target_latest"
    if [ "$legacy_available" = "yes" ] ; then
        sudo_commands+=" ; mv $tmpfile2 $target_legacy"
        if [ "$targetdir" != "/tmp" ] ; then
            sudo_commands+=" ; chown root:root $target_legacy"
        fi
        sudo_commands+=" ; chmod 0644 $target_legacy"
    else
        sudo_commands+=" ; if [ -r $target_legacy ] ; then mv -f $target_legacy $target_legacy.unsupported ; fi"
    fi
    #sudo_commands+=" ; echo $readme > $targetdir/README.txt"

    if [ "$targetdir" = "/tmp" ] ; then
        # temporary files used, no sudo needed
        Info "Adding NVIDIA graphics card ids to a temporary local database."
        bash -c "$sudo_commands"
    else
        # Now we need elevated privileges.
        Info "Adding NVIDIA graphics card ids to the local database."
        $EOS_ROOTER "$sudo_commands"
    fi
}

CreateNvidiaDb "$@"    # program starts here
