#!/bin/bash

# Find the right Nvidia driver for a card.

echo2()   { echo   "$@" >&2 ; }
printf2() { printf "$@" >&2 ; }

DIE() {
    echo2 "$progname: error: $1"
    exit 1
}

FilterIds() {
    # supports series 390 and later
    grep '<tr id="devid' | sed 's|.*devid\([0-9A-F]*\)["_].*|\1|' | sort -u | tr '[:upper:]' '[:lower:]'
}

GetNvidiaIds() {
    local series="$1"
    [ -n "$series" ] || return 1
    local ids=""
    local sid1="${series_ids[0]}"   # e.g. 520
    local sid2="${series_ids[1]}"   # e.g. 470
    local sid3="${series_ids[2]}"   # e.g. 390
    local sidX="$series_id_out"

    case "$series" in
        $sid1) ids="$(echo "$ids_all" | sed -n '1,/.*Below are the legacy GPUs that are no longer supported.*/p' | FilterIds)" ;;
        $sid2) ids="$(echo "$ids_all" | sed -n "/.*\"legacy_$sid2.xx\"*/,/.*\"legacy_$sid3.xx\".*/p" | FilterIds)" ;;
        $sid3) ids="$(echo "$ids_all" | sed -n "/.*\"legacy_$sid3.xx\"*/,/.*\"legacy_$sidX.xx\".*/p" | FilterIds)" ;;

        $beta_series) ids="$(echo "$ids_beta" | sed -n '1,/.*Below are the legacy GPUs that are no longer supported.*/p' | FilterIds)" ;;

        *) DIE "unsupported series '$series'" ;;
    esac

    echo "$ids"
}

ShowDriversForId() {
    local id="$1"
    local ix
    local ids
    local ver
    local true_series=""   # found from the Nvidia website
    local ids_all
    local ids_beta
    local found=no
    local exit_code=0
    local -r url_of_ids="$(UrlOfIds "$pacman_pkgver")"
    local -r url_of_beta_ids="$(UrlOfIds "$beta_pkgver")"

    ids_all="$(curl --fail -Lsm $curl_timeout -o- "$url_of_ids")"
    exit_code=$?

    [ -n "$ids_all" ] || DIE "Nvidia id data not available from '$url_of_ids'! Curl exit code $exit_code."

    ids_beta="$(curl --fail -Lsm $curl_timeout -o- "$url_of_beta_ids")"
    [ -n "$ids_beta" ] || DIE "Nvidia id data not available from '$url_of_beta_ids'! Curl exit code $exit_code."

    for ver in "${nvidia_versions[@]}" ; do
        true_series="${ver%%.*}"
        if [ "$testing" = "yes" ] ; then
            [ $true_series -eq 390 ] || continue
        fi
        ids="$(GetNvidiaIds "$true_series")"
        if [ -n "$(echo "$ids" | grep "$id")" ] ; then
            printf "\nSeries %s: supported (nvidia.com: %s)\n\n" "$true_series" "$ver"     # DO NOT CHANGE THIS LINE! Other apps use it.
            if [ "$ver" = "$beta_pkgver" ] && [ "$beta_pkgver" != "$pacman_pkgver" ] ; then
                printf2 "NOTE: $beta_pkg $ver is an AUR package. To install it, use command:\n   yay -S $beta_pkg\n"
            fi
            case "$true_series" in
                ${series_ids[1]} | ${series_ids[2]} | 470 | 390)
                    if [ $hide_nvidia_inst_msgs = no ] ; then
                        printf2 "NOTE: nvidia-inst does not support series $true_series. You can install the driver using 'yay'.\n"
                    fi
                    ;;
            esac
            [ "$all" = "no" ] && return
            found=yes
        else
            [ "$all" = "yes" ] && echo2 "Series $true_series: not supported"
        fi
    done
    if [ "$found" = "no" ] ; then
        echo2 "Sorry, none of the Nvidia driver series [$(echo "${series_ids[*]}" | sed 's| |, |g')] is supported."
        printf2 "\nMore info:\n   %s\n   %s\n" "$starting_url" "$url_of_ids"
    fi
}

GetNvidiaVersions() {
    echo2 "NVIDIA card id: $nvidia_id"
    echo2 "Fetching driver data from nvidia.com ..."

    local s n
    local ret=0
    local versions
    versions=$(curl --fail -Lsm $curl_timeout -o- "$starting_url")       # get the html data about the versions of the Nvidia drivers
    ret=$?
    [ -n "$versions" ] || DIE "Fetching Nvidia driver info from '$starting_url' failed (curl code $ret). Please try again later."

    versions=$(echo "$versions" | grep "Linux x86_64/AMD64")                   # extract the line with only Linux x86_64 versions
    versions=$(echo "$versions" | sed 's|</a>|</a>\n|g' | grep Latest)         # make each version appear on a separate line
    versions=$(echo "$versions" | sed -E 's|.*>[ ]*([0-9\.]+)[ ]*<.*|\1|')     # extract version number on each line between '>' and '<' (excluding spaces)

    for n in $versions ; do
        [ $(vercmp $n $series_id_out) -lt 0 ] && continue
        for s in "${series_ids[@]}" ; do
            case "$n" in
                $s.*) nvidia_versions+=("$n") ;;
            esac
        done
    done
    if [ "$beta_pkgver" != "$pacman_pkgver" ] ; then
        nvidia_versions+=("$beta_pkgver")
    fi
}

UrlOfIds() {
    local ver="$1"
    echo "https://download.nvidia.com/XFree86/Linux-x86_64/$ver/README/supportedchips.html"
    # echo "https://us.download.nvidia.com/XFree86/Linux-x86_64/$version/README/supportedchips.html"   # no more works (?)
}

GetPacmanVersion() {
    local infopkg="$1"
    {
        if [ -x /usr/bin/expac ] ; then
            expac -S %v "$infopkg"
        else
            LANG=C pacman -Si "$infopkg" | grep -w ^Version | awk '{print $NF}'
        fi
    } | head -n1   # in case multiple repos have it, take the first
}

Options() {
    local opts

    opts="$(/usr/bin/getopt -o=hajt: --longoptions help,all,test,timeout:,hide-nvidia-inst-msgs --name "$progname" -- "$@")" || {
        Options -h
        return 1
    }

    eval set -- "$opts"

    while true ; do
        case "$1" in
            -a | --all) echo2 "Info: option '$1' recognized but is no more used." ;; # all=yes ;;
            -j | --test) testing=yes ;;
            -t | --timeout) curl_timeout="$2" ;;
            --hide-nvidia-inst-msgs)
                # Do not show messages related to program nvidia-inst.
                hide_nvidia_inst_msgs=yes
                ;;

            -h | --help)
                cat <<EOF >&2

$progname finds the latest Nvidia driver series info that supports
your Nvidia graphics card. This information should help in installing the proper Nvidia driver.

Usage: $progname [options] [parameters]

Options:
  -t, --timeout     Timeout in seconds for operations that fetch driver data from nvidia.com.
                    Default: $curl_timeout_default
  -a, --all         (Show status for all available driver branches.)
                    NOTE: this option does nothing due to implementation change.

  -h, --help        This help.

Parameters:
  card_id           Nvidia card id (4 hex numbers, lower case letters).
                    Example: 1b83
                    Note: card_id is an optional parameter.

Tip: the "card_id" value can be found with command:

  lspci -vnn | grep -P 'VGA|3D|Display' | sed -E 's|.*\[10de:([0-9a-f]*)\].*|\1|'

See also: $starting_url
EOF
                exit 0
                ;;
            
            --) shift ; break ;;
        esac
        shift
    done

    nvidia_id="$1"

    case "$nvidia_id" in
        "") ;;
        [0-9A-Fa-f][0-9A-Fa-f][0-9A-Fa-f][0-9A-Fa-f]) nvidia_id=${nvidia_id,,} ;;
        *) DIE "given card id '$nvidia_id' is invalid, must contain only hexadecimal numbers (0-9 a-f A-F)" ;;
    esac
}

Main() {
    local -r progname="${0##*/}"
    local -r infopkg=nvidia-dkms
    local -r starting_url="https://www.nvidia.com/en-us/drivers/unix"
    local -r all=no                                                      # this will no more be changed!
    local testing=no
    local -r curl_timeout_default=30
    local curl_timeout=$curl_timeout_default
    local nvidia_id=""
    local nvidia_versions=()
    local hide_nvidia_inst_msgs=no
    local -r beta_pkg=nvidia-beta-dkms

    Options "$@"

    local -r pacman_version=$(GetPacmanVersion "$infopkg")
    local -r pacman_pkgver=$(echo "$pacman_version" | sed 's|-[0-9.]*$||')
    local -r pacman_version_series=${pacman_pkgver%%.*}

    local -r beta_pkgver=$(yay -Ss $beta_pkg | grep ^aur/ | awk '{print $2}' | sed 's|-[0-9.]*$||')
    local -r beta_series=${beta_pkgver%%.*}

    [ -n "$pacman_pkgver" ] || DIE "cannot determine the version of the latest available package $infopkg"
    [ -n "$beta_pkgver" ]   || DIE "cannot determine the version of the latest available package $beta_pkg"

    local series_ids=(
        $pacman_version_series        # latest provided by Arch
        470                           # legacy series ids
        390                           # ditto
    )
    local series_id_out=367           # may change as series_ids change!

    if [ -z "$nvidia_id" ] ; then
        nvidia_id="$(lspci -vnn | grep -P 'VGA|Display|3D' | grep "\[10de:" | sed 's|.*\[10de:\([0-9a-f]*\).*|\1|')"
    fi

    GetNvidiaVersions                 # find versions NVIDIA is currently supporting

    if [ -n "$nvidia_id" ] ; then
        ShowDriversForId "$nvidia_id"
    fi
}

Main "$@"
