#!/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="560.35.03-6"  # Update this to the latest version
PACKAGE=
RET=1

# Map driver versions to their respective .ids files and Arch package names
declare -A DRIVER_IDS=(
    [71]="nvidia-legacy-71xx"
    [96]="nvidia-legacy-96xx"
    [173]="nvidia-legacy-173xx"
    [304]="nvidia-legacy-304xx"
    [340]="nvidia-legacy-340xx"
    [390]="nvidia-legacy-390xx"
    [470]="nvidia-tesla-470"
    [open]="nvidia-open"
    [latest]="nvidia"
)

NV_DETECT() {
    NVGA=$1
    IDLISTDIR=/usr/share/nvidia
    local VERSIONS=()

    # Check which drivers the GPU is supported by
    for version in "${!DRIVER_IDS[@]}"; do
        if grep -q -i $NVGA "$IDLISTDIR/${DRIVER_IDS[$version]}.ids" 2>/dev/null; then
            VERSIONS+=("$version")
        fi
    done

    if [[ ${#VERSIONS[@]} == 0 ]]; then
        echo "Your card is not supported by any driver version up to $LATEST."
        echo "Check the Arch Wiki or install newer drivers from the AUR."
        return
    fi

    # Determine the most appropriate driver
    if [[ " ${VERSIONS[@]} " =~ "latest" ]]; then
        echo "Your card is supported by the latest drivers."
        PACKAGE="nvidia"
    elif [[ " ${VERSIONS[@]} " =~ "open" ]]; then
        echo "Your card is supported by the open-source drivers."
        PACKAGE="nvidia-open"
    elif [[ " ${VERSIONS[@]} " =~ "470" ]]; then
        echo "Your card is supported by the Tesla dkms drivers."
        PACKAGE="nvidia-470xx-dkms"
    elif [[ " ${VERSIONS[@]} " =~ "390" ]]; then
        echo "Your card is supported by the legacy 390xx drivers."
        PACKAGE="nvidia-390xx-dkms"
    elif [[ " ${VERSIONS[@]} " =~ "340" ]]; then
        echo "Your card is supported by the legacy 340xx drivers."
        PACKAGE="nvidia-340xx-dkms"
    elif [[ " ${VERSIONS[@]} " =~ "304" ]]; then
        echo "Your card is supported by the legacy 304xx drivers."
        PACKAGE="nvidia-304xx"
    elif [[ " ${VERSIONS[@]} " =~ "173" ]]; then
        echo "Your card is supported by the legacy 173xx drivers."
        PACKAGE="nvidia-173xx"
    elif [[ " ${VERSIONS[@]} " =~ "96" ]]; then
        echo "Your card is supported by the legacy 96xx drivers."
        PACKAGE="nvidia-96xx"
    elif [[ " ${VERSIONS[@]} " =~ "71" ]]; then
        echo "Your card is supported by the legacy 71xx drivers."
        PACKAGE="nvidia-71xx"
    else
        echo "Oops. Internal error ($NVGA)"
    fi

    if [ "$PACKAGE"  == "nvidia" ] && [ "$PACKAGE" == "nvidia-open" ]; then
        echo "It is recommended to install the nvidia package or install the nvidia-dkms package for custom kernels"
        echo "Also, If you want to have an open source driver with good performance, you can install the nvidia-open package or install the nvidia-open-dkms package for custom kernels"
        RET=0
    elif [ "$PACKAGE"  == "nvidia" ]; then
        echo "It is recommended to install the $PACKAGE package or install the nvidia-dkms package for custom kernels"
        RET=0
    elif [ "$PACKAGE" == "nvidia-304xx" ] || [ "$PACKAGE" == "nvidia-173xx" ] || [ "$PACKAGE" == "nvidia-96xx" ] || [ "$PACKAGE" == "nvidia-71xx" ]; then
        echo "Unfortunately, Arch Linux does not support $PACKAGE driver."
        echo "It is recommended to install the mesa package which is installed by default or install Debian distribution."
        RET=0
    elif [ -n "$PACKAGE" ]; then
        echo "It is recommended to install the $PACKAGE package from aur."
        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
