#!/bin/bash

# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: LicenseRef-NvidiaProprietary
#
# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
# property and proprietary rights in and to this material, related
# documentation and any modifications thereto. Any use, reproduction,
# disclosure or distribution of this material and related documentation
# without an express license agreement from NVIDIA CORPORATION or
# its affiliates is strictly prohibited.
#

# Publish the CUDA userspace library directory for the active GPU variant
# (same selection logic as nv-load-display-modules.sh) via
# ld.so.conf.d so the dynamic linker prefers either nvidia-l4t-cuda-nvgpu or
# nvidia-l4t-cuda-openrm under /opt/nvidia/l4t-gpu-libs/{nvgpu,openrm}/.
# Also publish /usr/lib/aarch64-linux-gnu/libcuda.so and libcuda.so.1 as
# symlinks to the active variant's libraries for link-time users (covers
# both `-lcuda` and `-l:libcuda.so.1` linker forms).
#
# Sources nv-load-display-modules-choose-variant.sh and maps variant names
# (e.g. nvgpu-l4t, openrm-l4t) to nvgpu or openrm by substring. Does not run
# ldconfig; ldconfig.service refreshes the cache after this unit runs.
#
# Safe to re-run. Removes /etc/ld.so.cache when the selected path changes so
# that ConditionFileNotEmpty=|!/etc/ld.so.cache causes ldconfig.service to
# rebuild the cache. On unknown variant or missing library directory: remove
# the ld.so fragment and CUDA linker symlink.

set -eo pipefail

CHOOSE_VARIANT_HELPER="/etc/systemd/nv-load-display-modules-choose-variant.sh"
# 000- prefix: process before other ld.so.conf.d fragments (e.g. nvidia-tegra) so
# the selected /opt/nvidia/l4t-gpu-libs path is added to the cache with precedence.
CONF="/etc/ld.so.conf.d/000-nvidia-gpu-libs.conf"
GPU_LIBS_ROOT="/opt/nvidia/l4t-gpu-libs"
CUDA_LINK_DIR="/usr/lib/aarch64-linux-gnu"
CUDA_LINK="${CUDA_LINK_DIR}/libcuda.so"
CUDA_LINK_SONAME="${CUDA_LINK_DIR}/libcuda.so.1"

if [ ! -f "${CHOOSE_VARIANT_HELPER}" ]; then
	echo "ERROR: nv-load-gpu-libs: required helper not found: ${CHOOSE_VARIANT_HELPER}" >&2
	exit 1
fi
. "${CHOOSE_VARIANT_HELPER}"

variant="$(nv_load_display_modules_choose_variant)"
v="$(printf '%s' "${variant}" | tr '[:upper:]' '[:lower:]')"

suffix=""
# Prefer nvgpu when both substrings could match (pathological input).
if [[ "${v}" == *nvgpu* ]]; then
	suffix="nvgpu"
elif [[ "${v}" == *openrm* ]]; then
	suffix="openrm"
fi

lib_dir=""
if [ -z "${suffix}" ]; then
	echo "ERROR: nv-load-gpu-libs: GPU variant is not nvgpu or openrm: ${variant}" >&2
elif [ ! -d "${GPU_LIBS_ROOT}/${suffix}" ]; then
	echo "ERROR: nv-load-gpu-libs: library directory missing: ${GPU_LIBS_ROOT}/${suffix}" >&2
else
	lib_dir="${GPU_LIBS_ROOT}/${suffix}"
fi

if [ -z "${lib_dir}" ]; then
	rm -f "${CUDA_LINK}" "${CUDA_LINK_SONAME}"
	if [ -f "${CONF}" ]; then
		rm -f "${CONF}"
		# Remove cache so ConditionFileNotEmpty=|!/etc/ld.so.cache
		# triggers ldconfig.service.
		rm -f /etc/ld.so.cache
	fi
else
	expected="$(printf '%s\n' "${lib_dir}")"
	if [ "$(cat "${CONF}" 2>/dev/null)" != "${expected}" ]; then
		mkdir -p /etc/ld.so.conf.d
		printf '%s\n' "${lib_dir}" > "${CONF}"
		mkdir -p "${CUDA_LINK_DIR}"
		ln -sfnT "${lib_dir}/libcuda.so" "${CUDA_LINK}"
		ln -sfnT "${lib_dir}/libcuda.so.1" "${CUDA_LINK_SONAME}"
		# Remove cache so ConditionFileNotEmpty=|!/etc/ld.so.cache
		# triggers ldconfig.service.
		rm -f /etc/ld.so.cache
	fi
fi
