#!/bin/bash

# Find the right Nvidia driver for a card.

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

DIE()  { echo2 "==> $progname: error: $1"; exit 1; }
INFO() { echo2 "==> $progname: info: $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)" ;;
    esac
    if [ -z "$ids" ] && [ $include_beta = yes ] ; then
        case "$series" in
            $beta_series) ids="$(echo "$ids_beta" | sed -n '1,/.*Below are the legacy GPUs that are no longer supported.*/p' | FilterIds)" ;;
        esac
    fi
    if [ "$ids" ] ; then
        echo "$ids"
        return
    fi

    DIE "unsupported series '$series'"
}

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

    if [ $include_beta = yes ] ; then
        local ids_beta
        local -r url_of_beta_ids="$(UrlOfIds "$beta_pkgver")"
    fi

    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."

    if [ $include_beta = yes ] ; then
        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."
    fi

    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 [ $include_beta = yes ] ; then
                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
            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 [ $include_beta = yes ] ; then
        if [ "$beta_pkgver" != "$pacman_pkgver" ] ; then
            nvidia_versions+=("$beta_pkgver")
        fi
    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
}

HasAurSupport() {
    local aur_timeout=10
    curl -Lsm $aur_timeout https://aur.archlinux.org/packages >/dev/null
}

Options() {
    local opts

    opts="$(/usr/bin/getopt -o=habjt: --longoptions help,all,beta,test,timeout:,hide-nvidia-inst-msgs --name "$progname" -- "$@")" || Help 1
    eval set -- "$opts"

    while true ; do
        case "$1" in
            -b | --beta)             include_beta=yes ;;
            -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) hide_nvidia_inst_msgs=yes ;;               # Do not show messages related to program nvidia-inst.
            -h | --help)             Help 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
}

Help() {
    cat <<EOF >&2

$progname finds info about the latest Nvidia driver series that supports
your Nvidia GPU. 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.
  -b, --beta        (Experimental) Include beta packages in the check if AUR is available.

  -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
    [ "$1" ] && exit $1
}

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 beta_pkg=nvidia-beta-dkms
    local include_beta=no               # check also beta package from the AUR if "yes"; "no" does not check beta

    Options "$@"

    if [ $include_beta = yes ] ; then
        HasAurSupport || { include_beta=no; INFO "sorry, AUR is not currently available, cannot include beta packages"; }
    fi

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

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

    if [ $include_beta = yes ] ; then
        local beta_pkgver=$(LANG=C /bin/yay -Si $beta_pkg | grep -w "^Version" | awk '{print $NF}')
        beta_pkgver=${beta_pkgver%-*}
        # local beta_pkgver=$(yay -Ss $beta_pkg | grep "aur/${beta_pkg}" | awk '{print $2}' | sed 's|-[0-9.]*$||')
        local -r beta_series=${beta_pkgver%%.*}
        [ -n "$beta_pkgver" ]   || DIE "cannot determine the version of the latest available package $beta_pkg"
    fi

    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 "$@"
