#!/usr/bin/env nix-shell
#!nix-shell -i bash -p bash coreutils gnugrep gum
#
# Based on https://wiki.debian.org/NvidiaGraphicsDrivers?action=AttachFile&do=view&target=nvidia-versions.sh
#
# Copyright © 2026 Raphaël Numbus <gpu-detect@numbus.eu>
#
# This package is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 3 of the License, or
# (at your option) any later version.
#
# This package is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program. If not, see <https://www.gnu.org/licenses/>

gum style --border-foreground="030" --bold --border=normal --padding="0 2" "GPU detection tool for NixOS"

USER_ID="$(id -u)"

declare -A args=(
  ["brand"]="all"
  ["hw_accel"]="true"
)

while [[ $# -gt 0 ]]; do
  case "${1}" in
    --brand|-b)
      args["brand"]="${2}"
      shift 2
      ;;
    --hw_accel|-hw)
      args["hw_accel"]="${2}"
      shift 2
      ;;
    --help|-h|*)
      echo "Usage: gpu-detect --brand=[all/amd/intel/nvidia] --hw_accel=[true/false] [PCIID]..."
      echo ""
      echo "For each GPU installed in your system, the program"
      echo "will provide you the needed NixOS configuration."
      exit 0
      ;;
  esac
done

shopt -s compat31 nocasematch 2>/dev/null || { echo "Error: this script only works with bash." && exit; } # Avoid cryptic failure when running dash on this script

# last time the PCI IDs were updated
LATEST="#VERSION#"
PACKAGE=
RET=1

amd_detect() {
	DETECT_LOCATION="/run/user/${USER_ID}/gpu-detect/amd"
	mkdir -p "${DETECT_LOCATION}"

	curl -sL https://gittea.dev/numbus/gpu-detect/raw/branch/26.05/server/ids-database/nvidia/legacy/legacy_340.ids -o "${DETECT_LOCATION}"/legacy_340.ids
	curl -sL https://gittea.dev/numbus/gpu-detect/raw/branch/26.05/server/ids-database/nvidia/legacy/legacy_390.ids -o "${DETECT_LOCATION}"/legacy_390.ids
	curl -sL https://gittea.dev/numbus/gpu-detect/raw/branch/26.05/server/ids-database/nvidia/legacy/legacy_470.ids -o "${DETECT_LOCATION}"/legacy_470.ids
	curl -sL https://gittea.dev/numbus/gpu-detect/raw/branch/26.05/server/ids-database/nvidia/legacy/legacy_580.ids -o "${DETECT_LOCATION}"/legacy_580.ids
}

intel_detect() {

}

nvidia_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-open.ids 2>/dev/null
then
	VERSIONS[990]=515.48.07
fi

if grep -q -i $NVGA $IDLISTDIR/nvidia.ids 2>/dev/null
then
	# 999 means current
	VERSIONS[999]=$LATEST
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 "Newer driver releases may be available in backports, unstable or experimental."
	return
fi



if [ -n "$PACKAGE" ]; then
	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