#
# Copyright 2024 Google LLC
#
# This program is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License
# version 2 as published by the Free Software Foundation.
#
# This program 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.
#

EAPI=7

inherit cos-linux-fw cos-linux-mod linux-mod eutils toolchain-funcs unpacker

DESCRIPTION="NVIDIA Accelerated Graphics Driver Precompilation"

NV_KERNEL_MAX="6.4"
NV_URI="https://us.download.nvidia.com/tesla/"

INSTALLER_NAME="NVIDIA-Linux-x86_64-${PV}"
RUNFILE="${INSTALLER_NAME}.run"
PRECOMPILED="${INSTALLER_NAME}-custom"
DESCRIPTION="NVIDIA Accelerated Graphics Driver"
HOMEPAGE="https://www.nvidia.com/download/index.aspx"
SRC_URI="${NV_URI}${PV}/${RUNFILE}"

S="${WORKDIR}"

LICENSE="NVIDIA-r2 BSD BSD-2 GPL-2 MIT ZLIB curl openssl"
SLOT="0/${PV%%.*}"
KEYWORDS="*"
IUSE="+clang kernel-5_10 kernel-5_15"

COMMON=""
DEPEND="
	${COMMON}
	kernel-5_10? ( sys-kernel/lakitu-kernel-5_10[kernel_sources] )
	kernel-5_15? ( sys-kernel/lakitu-kernel-5_15[kernel_sources] )
"
RDEPEND="
	${COMMON}"

src_unpack() {
	# we want the entire .run file
	# using the default src_unpack extracts the runfile contents.
	cp ${DISTDIR}/${RUNFILE} ${S}
}

src_prepare() {
	use clang || cros_use_gcc
	default
}

# create a wrapper that removes -Werror=strict-prototypes
# from the compiler options
create_cc_wrapper() {
	cp "${FILESDIR}/nvidia-installer-cc.wrapper" "${S}/${CC}"
	chmod a+x "${S}/${CC}"
}

# the $LD points to the ld.lld linker which has a bug
# .bss relocations are not done properly with the nv-kernel.ld linker script.
# ld.bfd doesn't seem to have this issue, but we don't use ld.bfd to compile our kernel.
force_symlink_bfd_ld() {
	ln -sf "/usr/bin/${CROSS_COMPILE}ld.bfd" "${S}/ld"
}

src_compile() {
	cros_allow_gnu_build_tools
	local -r kernel_headers_tarball="${SYSROOT}/opt/google/src/kernel-headers.tgz"
	tar xf "${kernel_headers_tarball}" -C "${S}"
	local -r kernel_headers_path="${S}/usr/src/linux-headers-${KV_FULL}"
	create_cc_wrapper
	force_symlink_bfd_ld
	local -r logfile="${S}/nvidia-installer.log"
	sudo PATH="${S}:${PATH}" \
		CC="${CC}" \
		sh ${RUNFILE} \
		--kernel-source-path="${kernel_headers_path}" \
		--add-this-kernel \
		--no-install-compat32-libs --silent --accept-license  \
		--log-file-name="${logfile}" || die "unable to produce precompiled drivers runfile"
}

src_install() {
	# extracting precompiled drivers
	sh "${PRECOMPILED}.run" --extract-only
	precompiled_archive="$(find ${PRECOMPILED}/kernel/precompiled -name nvidia-precompiled*)"
	echo "precompiling......................."
	sudo "./${PRECOMPILED}/mkprecompiled" \
		"--unpack" \
		"$precompiled_archive" \
		"-o" \
		"${PRECOMPILED}/kernel/precompiled"

	# link
	linkerscript="${S}/usr/src/linux-headers-${KV_FULL}/scripts/module.lds"
	# link nvidia.ko
	$(tc-getLD) -T ${linkerscript} \
		-r -o "${PRECOMPILED}/kernel/precompiled/nvidia.ko" \
		"${PRECOMPILED}/kernel/precompiled/nv-linux.o" \
		"${PRECOMPILED}/kernel/nvidia/nv-kernel.o_binary"

	# link nvidia-modeset.ko
	$(tc-getLD) -T ${linkerscript} \
		-r -o "${PRECOMPILED}/kernel/precompiled/nvidia-modeset.ko" \
		"${PRECOMPILED}/kernel/precompiled/nv-modeset-linux.o" \
		"${PRECOMPILED}/kernel/nvidia-modeset/nv-modeset-kernel.o_binary"
	if use module_sign ; then
		TMP_DIR=$(mktemp -d)
		SIGNATURES_DIR="${TMP_DIR}/signatures"
		mkdir -p "${SIGNATURES_DIR}"
		MODULE_OBJECTS=(nvidia.ko nvidia-drm.ko nvidia-modeset.ko nvidia-peermem.ko nvidia-uvm.ko)
		pushd "${PRECOMPILED}/kernel/precompiled/" || die
		for MODULE_OBJECT in "${MODULE_OBJECTS[@]}"; do
			# Sign the nvidia modules.
			cos-linux-mod_sign_detached "${MODULE_OBJECT}" "${KBUILD_OUTPUT}/certs" || die "signing failed"
			# copy the object file with extended attributes.
			cp "${MODULE_OBJECT}.p7s" "${SIGNATURES_DIR}/${MODULE_OBJECT}.sig"
		done
		popd || die

		FIRMWARE_OBJECTS=(gsp_ga10x.bin gsp_tu10x.bin)
		pushd "${PRECOMPILED}/firmware" || die
		for FIRMWARE_OBJECT in "${FIRMWARE_OBJECTS[@]}"; do
			# Sign GSP firmware modules.
			cos-linux-fw_sign_detached "${FIRMWARE_OBJECT}" || die "signing failed"
			# copy the object file with extended attributes.
			cp --preserve=all "${FIRMWARE_OBJECT}.sig" "${SIGNATURES_DIR}"/
		done
		popd || die
		# Nvidia drivers are stored as build artifacts and not installed in image.
		local source_dir=opt/google/drivers
		dodir ${source_dir}
		# tar the files with extended attributes.
		tar -czf "${D}/${source_dir}/nvidia-drivers-${PV}-signature.tar.gz" -C "${SIGNATURES_DIR}"/ .  || die
		rm -rf "${TMP_DIR}"
	fi

	# precompiled driver stored as build artifacts
	# exported out of the image
	local source_dir=opt/google/drivers
	dodir ${source_dir}
	mv "${S}/${INSTALLER_NAME}-custom.run" "${D}/${source_dir}"
}