#!/usr/bin/env nix-shell
#!nix-shell -i bash -p bash coreutils gnugrep gum git pciutils
#
# 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/>

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

# Defining variables
USER_ID="$(id -u)"
PCI_IDS_LOCATION="/run/user/${USER_ID}/gpu-detect"
declare -A args=(
  ["brand"]="all"
  ["hw_accel"]="true"
)


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



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

gpu_detect() {
	# Get the PCI IDs from the repo
	if [[ ! -d "${PCI_IDS_LOCATION}" ]]; then
		mkdir -p "${PCI_IDS_LOCATION}"
		git clone -q -b ids-database https://gittea.dev/numbus/gpu-detect.git "${PCI_IDS_LOCATION}"
	else
		if ! git -C "${PCI_IDS_LOCATION}" pull; then
			echo "Directory ${PCI_IDS_LOCATION} is not empty and does not"
			echo "contain the repository. Please check the folder clean up manually."
			exit 1
		fi
	fi

	amd_detect() {

	}

	intel_detect() {

	}

	nvidia_detect() {
		PCI_ID=${1}
		NVIDIA_VERSION=""
		local NVIDIA_TEXT="""
			{ config, lib, pkgs, ... }:
			{
				hardware.graphics = {
					enable = true;
					extraPackages = with pkgs; [ ${EXTRA_PACKAGES} ];
				};
				services.xserver.videoDrivers = [ "nvidia" ];
				hardware.nvidia = {
					modesetting.enable = true;
					nvidiaSettings = true;
					package = config.boot.kernelPackages.nvidiaPackages.${NVIDIA_VERSION};

					# Unsupported options
					powerManagement.enable = false;
					powerManagement.finegrained = false;
					open = false;
				};
			}
		"""
		local VERSIONS=("unsupported" "legacy_340" "legacy_390" "legacy_470" "legacy_580" "current")

		for version in ${VERSIONS[@]}; do
			if grep -q -i "${PCI_ID}" "${PCI_IDS_LOCATION}/nvidia/${version}.ids" 2>/dev/null; then
				NVIDIA_VERSION="${version}"
				break
			fi
		done

		if [[ -z "${NVIDIA_VERSION}" ]]; then
			echo "Uh oh. It seems that your card couldn't be detected."
			echo "A newer driver may add support for your card later."
			return 1
		fi

		if [[ "${NVIDIA_VERSION}" == "unsupported" ]]; then
			echo "Good news! You have nothing to do, your card runs with the"
			echo "\"nouveau\" driver that is integrated in the kernel."
			echo ""
			echo "It seems you have a very old card!"
			echo "The proprietary NVIDIA drivers have been removed from Nixpkgs"
			echo "therefore it is not possible to install them."
		elif [[ "${NVIDIA_VERSION}" == "legacy_340" || "${NVIDIA_VERSION}" == "legacy_390" || "${NVIDIA_VERSION}" == "legacy_470" ]]; then
			echo "An official legacy NVIDIA driver release is available for"
			echo "your card."
			echo ""
			echo "You can add the following code snippet to your NixOS configuration"
			echo "in order to enable the proprietary NVIDIA driver with VDPAU support."
			EXTRA_PACKAGES="\"libvdpau\""
			echo ${NVIDIA_TEXT}
		elif [[ "${NVIDIA_VERSION}" == "legacy_580" ]]; then
			echo "An official legacy NVIDIA driver release is available for"
			echo "your card."
			echo ""
			echo "You can add the following nix code to your NixOS configuration"
			echo "in order to enable the proprietary NVIDIA driver with VDPAU and"
			echo "VAAPI (if applicable) support."
			EXTRA_PACKAGES="\"nvidia-vaapi-driver\" \"libvdpau\""
			echo ${NVIDIA_TEXT}
		elif [[ "${NVIDIA_VERSION}" == "current" ]]; then
			echo "An official stable NVIDIA driver release is available for"
			echo "your card."
			echo ""
			echo "You can add the following nix code to your NixOS configuration"
			echo "in order to enable the proprietary NVIDIA driver with hardware"
			echo "acceleration support."
			echo """
				{ config, lib, pkgs, ... }:
				{
					hardware.graphics = {
						enable = true;
						extraPackages = with pkgs; [ ${EXTRA_PACKAGES} ];
					};
					services.xserver.videoDrivers = [ "nvidia" ];
					hardware.nvidia = {
						modesetting.enable = true;
						nvidiaSettings = true;
						open = true;
						package = config.boot.kernelPackages.nvidiaPackages.${NVIDIA_VERSION};

						# Advanced options, tweak if needed
						powerManagement.enable = false; # Experimental: Enable this if you have graphical corruption issues or application crashes after waking up from sleep.
						powerManagement.finegrained = false; # Experimental: Turns off GPU when not in use
					};
				}
			"""
		fi
	}


		if [ -z "${1}" ]; then


			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}')"
				nvidia_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/')]"
				nvidia_detect "$PCIID"
			done

		fi


}




exit $RET