#!/bin/bash
if [ "$1" = "-h" -o "$1" = "--help" ]; then
    echo "Usage: nvidia-detect [PCIID]..."
    echo " Reports the Arch packages supporting the NVIDIA GPU that is"
    echo " installed on the local system (or given as a PCIID parameter)."
    exit 0
fi
shopt -s compat31 nocasematch 2>/dev/null || { echo "Error: this script only works with bash." && exit; }
LATEST="590.00.00"
PACKAGE=
RET=1
NV_DETECT() {
    NVGA=$1
    IDLISTDIR=/usr/share/nvidia
    local VERSIONS
    if grep -q -i $NVGA $IDLISTDIR/nvidia-legacy-71xx.ids 2>/dev/null
    then
        VERSIONS[71]=71.86
    fi
    if grep -q -i $NVGA $IDLISTDIR/nvidia-legacy-96xx.ids 2>/dev/null
    then
        VERSIONS[96]=96.43
    fi
    if grep -q -i $NVGA $IDLISTDIR/nvidia-legacy-173xx.ids 2>/dev/null
    then
        VERSIONS[173]=173.14
    fi
    if grep -q -i $NVGA $IDLISTDIR/nvidia-legacy-304xx.ids 2>/dev/null
    then
        VERSIONS[304]=304.123
    fi
    if grep -q -i $NVGA $IDLISTDIR/nvidia-legacy-340xx.ids 2>/dev/null
    then
        VERSIONS[340]=340.76
    fi
    if grep -q -i $NVGA $IDLISTDIR/nvidia-legacy-390xx.ids 2>/dev/null
    then
        VERSIONS[390]=390.87
    fi
    if grep -q -i $NVGA $IDLISTDIR/nvidia-legacy-390xx-amd64.ids 2>/dev/null
    then
        VERSIONS[391]=390.87
    fi
    if grep -q -i $NVGA $IDLISTDIR/nvidia-tesla-418.ids 2>/dev/null
    then
        VERSIONS[418]=418.87.01
    fi
    if grep -q -i $NVGA $IDLISTDIR/nvidia-tesla-470.ids 2>/dev/null
    then
        VERSIONS[470]=470.57.02
    fi
    if grep -q -i $NVGA $IDLISTDIR/nvidia-tesla-535.ids 2>/dev/null
    then
        VERSIONS[535]=535.216.01
    fi
    if grep -q -i $NVGA $IDLISTDIR/nvidia-legacy-580xx.ids 2>/dev/null
    then
        VERSIONS[580]=580.00.00
    fi
    if grep -q -i $NVGA $IDLISTDIR/nvidia-open.ids 2>/dev/null
    then
        VERSIONS[990]=590.00.00
    fi
    if grep -q -i $NVGA $IDLISTDIR/nvidia.ids 2>/dev/null
    then
        if [[ -n ${VERSIONS[580]} ]]; then
            VERSIONS[999]=
        else
            VERSIONS[999]=$LATEST
        fi
    fi
    if [[ ${#VERSIONS[*]} == 0 ]]; then
        echo "Uh oh. Your card is not supported by any driver version up to $LATEST."
        echo "A newer driver may add support for your card."
        echo "Check the Arch Wiki or AUR for newer drivers."
        return
    fi
    if [[ -n ${VERSIONS[999]} ]]; then
        if [[ -n ${VERSIONS[535]} ]]; then
            echo "Your card is supported by all driver versions."
        else
            echo "Your card is supported by the latest drivers."
        fi
        echo ""
        echo "IMPORTANT: As of driver 590, Arch has switched to Open Kernel Modules."
        PACKAGE="nvidia-open"
    elif [[ -n ${VERSIONS[580]} ]]; then
        echo "Your card is Pascal (GTX 10xx series) or older."
        echo ""
        echo "IMPORTANT: Driver 590+ no longer supports Pascal and older cards."
        echo "You must use the legacy 580 driver series."
        echo ""
        echo "If you have nvidia/nvidia-dkms/nvidia-lts installed, you need to:"
        echo " 1. Uninstall the official nvidia packages"
        echo " 2. Install nvidia-580xx-dkms from AUR"
        PACKAGE="nvidia-580xx-dkms (AUR)"
    elif [[ -n ${VERSIONS[535]} ]]; then
        echo "Your card is supported by the Tesla 535 drivers series."
        PACKAGE="nvidia-535xx-dkms (AUR)"
    elif [[ -n ${VERSIONS[470]} ]]; then
        echo "Your card is supported by the Tesla 470 drivers series."
        PACKAGE="nvidia-470xx-dkms (AUR)"
    elif [[ -n ${VERSIONS[390]} ]]; then
        echo "Your card is supported by the legacy 390xx drivers series."
        PACKAGE="nvidia-390xx-dkms (AUR)"
    elif [[ -n ${VERSIONS[391]} ]]; then
        echo "Your card is only supported on the amd64 platform."
        echo "Your card is only supported up to the 390 legacy drivers series."
        PACKAGE="nvidia-390xx-dkms (AUR)"
    elif [[ -n ${VERSIONS[418]} ]]; then
        echo "Your card is supported by the Tesla 418 drivers series."
        PACKAGE="nvidia-418xx-dkms (AUR)"
    elif [[ -n ${VERSIONS[340]} ]]; then
        echo "Your card is only supported by the legacy 340xx drivers series."
        PACKAGE="nvidia-340xx-dkms (AUR)"
    elif [[ -n ${VERSIONS[304]} ]]; then
        echo "Uh oh. Your card is only supported by the 304 legacy drivers series."
        echo "This driver is very old and may not be available in AUR."
        PACKAGE="nvidia-304xx (check AUR)"
    elif [[ -n ${VERSIONS[173]} ]]; then
        echo "Uh oh. Your card is only supported by the 173.14 legacy drivers series."
        echo "This driver is very old and not available in Arch repositories."
    elif [[ -n ${VERSIONS[96]} ]]; then
        echo "Uh oh. Your card is only supported by the 96.43 legacy drivers series."
        echo "This driver is very old and not available in Arch repositories."
    elif [[ -n ${VERSIONS[71]} ]]; then
        echo "Uh oh. Your card is only supported by the 71.86 legacy drivers series."
        echo "This driver is very old and not available in Arch repositories."
    else
        echo "Oops. Internal error ($NVGA)"
    fi
    if [ -n "$PACKAGE" ] && [ "$PACKAGE" = "nvidia-open" ] && [[ -n ${VERSIONS[990]} ]]; then
        echo ""
        echo "Alternative packages for custom kernels:"
        echo " - nvidia-open-dkms (for custom kernels)"
        echo " - nvidia-lts-open (for linux-lts kernel)"
    fi
    if [ -n "$PACKAGE" ] && [ "$PACKAGE" != "nvidia-535xx-dkms (AUR)" ] && [[ -n ${VERSIONS[535]} ]]; then
        echo "Your card is also supported by the Tesla 535 drivers series."
    fi
    if [ -n "$PACKAGE" ]; then
        echo ""
        echo "It is recommended to install the"
        echo " $PACKAGE"
        echo "package."
        RET=0
    fi
}
if [ -z "$1" ]; then
    if ! (lspci --version) > /dev/null 2>&1; then
        echo "ERROR: The 'lspci' command was not found. Please install the 'pciutils' package." >&2
        exit 1
    fi
    NV_DEVICES=$(lspci -mn | awk '{ gsub("\"",""); if (($2 ~ "030[0-2]") && ($3 == "10de" || $3 == "12d2")) { print $1 } }')
    if [ -z "$NV_DEVICES" ]; then
        echo "No NVIDIA GPU detected."
        exit 0
    fi
    echo "Detected NVIDIA GPUs:"
    for d in $NV_DEVICES ; do
        lspci -nn -s $d
    done
    for d in $NV_DEVICES ; do
        echo -e "\nChecking card: $(lspci -s $d | awk -F: '{print $3}')"
        NV_DETECT "$(lspci -mn -s "$d" | awk '{ gsub("\"",""); print $3 $4 }')"
    done
else
    for id in "$@" ; do
        PCIID=$(echo "$id" | sed -rn 's/^(10de)?:?([0-9a-fA-F]{4})$/10de\2/ip')
        if [ -z "$PCIID" ]; then
            echo "Error parsing PCI ID '$id'."
            exit 2
        fi
        echo "Checking driver support for PCI ID [$(echo $PCIID | sed -r 's/(....)(....)/\1:\2/')]"
        NV_DETECT "$PCIID"
    done
fi
exit $RET
