#!/bin/bash

# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice, this
# list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright notice,
# this list of conditions and the following disclaimer in the documentation
# and/or other materials provided with the distribution.
#
# 3. Neither the name of the copyright holder nor the names of its
# contributors may be used to endorse or promote products derived from
# this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#

# Select the display driver variant (nvgpu-l4t or openrm-l4t), copy the appropriate
# depmod.d and modprobe.d conf files from
# /opt/nvidia/nv-disp-module-configs/, run depmod -a if the conf
# changed, then load the appropriate modules via modprobe.
#
# Safe to re-run: unloads any loaded display stack first (reverse order), then
# loads the chosen variant. Allows switching between nvgpu-l4t and openrm-l4t by
# updating /etc/nvidia-gpu-driver-override and re-running.
#
# /etc/nvidia-gpu-driver-override: optional override; write the desired
#   variant name (e.g. "nvgpu-l4t" or "openrm-l4t"). Otherwise: tegra23x
#   chips use nvgpu-l4t; tegra264 and later use openrm-l4t.
#
# /run/nvidia-gpu-driver-active: written unconditionally on each run;
#   contains the resolved variant name in effect for this boot.
#
# Override the depmod.d or modprobe.d conf file used:
#   /etc/nvidia-display-depmod-conf: if present, must contain the full path to
#     a depmod.d config file. Overrides the default packaged config.
#   /etc/nvidia-display-modprobe-conf: if present, must contain the full path
#     to a modprobe.d config file. Overrides the default packaged config.
#   These overrides are useful when using drivers installed by a method other
#   than the default NVIDIA packages (e.g. an NVIDIA installer or DKMS build)
#   that requires different module options or search paths.
#
# Configuration is cached in /var/lib/nvidia/nv-display.stamp. depmod -a and
# conf file copy are skipped when the selected config files, kernel version,
# and nvidia-l4t-init package version are all unchanged.
#
# openrm-l4t: modprobe nvidia_drm loads nvidia_modeset -> nvidia via modules.dep;
#   nv_gpu_static_pg (softdep pre:) and nvidia_uvm (softdep post:) are loaded
#   when nvidia is pulled in.
# nvgpu-l4t: modprobe nvgpu first (claims PCIe device), then modprobe nvidia_drm.

set -e

DISPLAY_DEPMOD_OVERRIDE="/etc/nvidia-display-depmod-conf"
DISPLAY_MODPROBE_OVERRIDE="/etc/nvidia-display-modprobe-conf"
KVER="$(uname -r)"
PKG_VER="$(dpkg-query -W -f='${Version}' nvidia-l4t-init 2>/dev/null || true)"
CONFIG_DIR="/opt/nvidia/nv-disp-module-configs"

DEPMOD_CONF="/etc/depmod.d/nvidia-display.conf"
MODPROBE_CONF="/etc/modprobe.d/nvidia-display.conf"
STAMP_FILE="/var/lib/nvidia/nv-display.stamp"
CHOOSE_VARIANT_HELPER="/etc/systemd/nv-load-display-modules-choose-variant.sh"

# Copy the resolved depmod.d and modprobe.d conf files to their system
# locations.
copy_conf_files() {
	local src_depmod="${1}"
	local src_modprobe="${2}"

	if [ ! -f "${src_depmod}" ]; then
		echo "ERROR: depmod config not found: ${src_depmod}" >&2
		exit 1
	fi
	if [ ! -f "${src_modprobe}" ]; then
		echo "ERROR: modprobe config not found: ${src_modprobe}" >&2
		exit 1
	fi

	mkdir -p /etc/depmod.d /etc/modprobe.d
	echo "Copying depmod config: ${src_depmod} -> ${DEPMOD_CONF}."
	cp "${src_depmod}" "${DEPMOD_CONF}"
	echo "Copying modprobe config: ${src_modprobe} -> ${MODPROBE_CONF}."
	cp "${src_modprobe}" "${MODPROBE_CONF}"
}

# Check if any graphics or display module is currently loaded.
is_graphics_display_loaded() {
	[ -d /sys/module/nvidia ] || [ -d /sys/module/nvgpu ]
}

# Unload graphics and display stack in reverse dependency order.
# Stderr is not hidden -- if you see "Module X is in use", something
# (e.g. display server) has the GPU. Stop that process and re-run.
unload_graphics_display_modules() {
	modprobe -r nvidia_uvm 2>/dev/null || true
	modprobe -r nvidia_drm 2>/dev/null || true
	modprobe -r nvidia 2>/dev/null || true
	modprobe -r nvgpu 2>/dev/null || true
}

if [ ! -f "${CHOOSE_VARIANT_HELPER}" ]; then
	echo "ERROR: required helper not found: ${CHOOSE_VARIANT_HELPER}" >&2
	exit 1
fi
. "${CHOOSE_VARIANT_HELPER}"
variant="$(nv_load_display_modules_choose_variant)"
echo "${variant}" > /run/nvidia-gpu-driver-active

# Resolve the depmod.d and modprobe.d config files for the selected variant.
# The defaults may be overridden by writing a full file path into
# /etc/nvidia-display-depmod-conf or /etc/nvidia-display-modprobe-conf.
src_depmod="${CONFIG_DIR}/nv-depmod-${variant}-display.conf"
src_modprobe="${CONFIG_DIR}/nv-modprobe-${variant}-display.conf"
if [ -f "${DISPLAY_DEPMOD_OVERRIDE}" ]; then
	override="$(tr -d '\n' \
		< "${DISPLAY_DEPMOD_OVERRIDE}" 2>/dev/null)"
	if [ -n "${override}" ]; then
		echo "Overriding depmod config: ${override}."
		src_depmod="${override}"
	fi
fi
if [ -f "${DISPLAY_MODPROBE_OVERRIDE}" ]; then
	override="$(tr -d '\n' \
		< "${DISPLAY_MODPROBE_OVERRIDE}" 2>/dev/null)"
	if [ -n "${override}" ]; then
		echo "Overriding modprobe config: ${override}."
		src_modprobe="${override}"
	fi
fi

# Copy conf files and run depmod -a only when the config changes.
desired_stamp="${src_depmod}:${src_modprobe}:${KVER}:${PKG_VER}"
current_stamp="$(cat "${STAMP_FILE}" 2>/dev/null)" || true
if [ "${desired_stamp}" != "${current_stamp}" ] || \
		[ ! -f "${DEPMOD_CONF}" ] || [ ! -f "${MODPROBE_CONF}" ]; then
	copy_conf_files "${src_depmod}" "${src_modprobe}"
	depmod -a
	mkdir -p "$(dirname "${STAMP_FILE}")"
	echo "${desired_stamp}" > "${STAMP_FILE}"
fi

# Unload only if something is already loaded.
if is_graphics_display_loaded; then
	unload_graphics_display_modules
fi

# For nvgpu-l4t, load nvgpu first so it claims the PCIe device before nvidia.ko
# registers its PCI driver; otherwise nvidia.ko probes the device and returns
# -EPERM (it is not the intended driver in this variant), producing a spurious
# kernel error message.
if [[ "${variant}" == *nvgpu* ]]; then
	modprobe nvgpu
fi

# Load display stack. nvidia_drm depends on nvidia_modeset and nvidia via
# modules.dep so modprobe nvidia_drm loads the full chain automatically.
# On openrm-l4t, nv_gpu_static_pg (softdep pre:) and nvidia_uvm (softdep post:)
# are also loaded when nvidia is pulled in as a dependency.
modprobe nvidia_drm

echo "Loaded display stack (variant=${variant}," \
	"depmod=${src_depmod}," \
	"modprobe=${src_modprobe})."
