#!/bin/bash

# SPDX-FileCopyrightText: Copyright (c) 2025 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.

set -e

BOOT_CTRL_CONF="/etc/nv_boot_control.conf"
OTA_PACKAGE_DIR="/opt/ota_package"
NVBOOTCTRL="/usr/sbin/nvbootctrl"
NV_PAYLOAD_UPDATER="/usr/sbin/nv_bootloader_payload_updater"
UPDATE_CONF="/opt/nvidia/l4t-bootloader-config/nv-l4t-bootloader-config.sh"
kernel_only_payload=""
OVERLAYFS_IS_ENABLED=0

if [ -e "/opt/nvidia/l4t-packages/.nv-l4t-disable-boot-fw-update-in-preinstall" ]; then
    echo "Pre-installing kernel package, skip flashing"
    exit 0
fi

# Select kernel BUP payloads based on the compatible spec and tnspec
compare_specs_to_select_payloads () {
    # Default payload
    kernel_only_payload="kernel_only_payload"

    # BUP payload selection only applies for t23x
    if [ "${chipid}" != "0x23" ]; then
        return 0
    fi

    # The compatible spec list for the AGX Orin boards with sku 0000
    # and with fab TSx/EBx/000/100/200
    local orin_boards_with_zero_fab=(
        '3701-000-0000--1--jetson-agx-orin-devkit-'
        '3701-000-0000--1--jetson-agx-orin-devkit-maxn-'
    )
    # shellcheck disable=SC2155
    local compatible_spec=$( awk '/COMPATIBLE_SPEC/ {print $2}' "${BOOT_CTRL_CONF}" )
    local found="false"

    if [ "${compatible_spec}" = "" ]; then
        echo "ERROR. Cannot find COMPATIBLE_SPEC in ${BOOT_CTRL_CONF}."
        return 1
    fi

    # Compare current board's compatible spec with the list
    for spec in "${orin_boards_with_zero_fab[@]}"; do
        if [ "${compatible_spec}" = "${spec}" ]; then
            found="true"
            break
        fi
    done
    if [ "${found}" != "true" ]; then
        return 0
    fi

    # Get the TNSPEC and parse board_fab in it
    # shellcheck disable=SC2155
    local tn_spec=$( awk '/TNSPEC/ {print $2}' "${BOOT_CTRL_CONF}" )
    if [ "${tn_spec}" = "" ]; then
        echo "ERROR. Cannot find TNSPEC in ${BOOT_CTRL_CONF}."
        return 1
    fi
    # shellcheck disable=SC2155
    local board_fab=$( echo "${tn_spec}" | awk -F"-" '{print $2}' )
    # Check the fabs that generated the zero fab filed in compatible spec.
    if [[ "${board_fab}" == "000" || "${board_fab}" == "TS"* || "${board_fab}" == "EB"*
       || ($((board_fab)) -gt 0 && $((board_fab)) -lt 300) ]]; then
        # Update the payload file if board FAB matches
        echo "INFO. Board fab matches, change kernel payload file."
        kernel_only_payload="kernel_only_payload_3701_000"
    else
        echo "INFO. ${board_fab} does not match any known boards."
    fi

    return 0
}

install_kernel () {
    part_name="${1}"

    "${NV_PAYLOAD_UPDATER}" --part "${part_name}" &> \
        "${OTA_PACKAGE_DIR}/${part_name}.log"
    # shellcheck disable=SC2181
    if [ $? -ne 0 ]; then
            echo "ERROR. Procedure for ${part_name} update FAILED."
            echo "Cannot install package. Exiting..."
            exit 1
    fi
}

# Install kernel and kernel-dtb to specified slot
install_kernel_on_rootfs_id () {
    local rootfs_id="${1}"
    local part_kernel=""
    local part_dtb=""

    if [ -x "${NV_PAYLOAD_UPDATER}" ]; then
        if [ "${rootfs_id}" = "1" ]; then
            part_kernel="B_kernel"
            part_dtb="B_kernel-dtb"
        else
            part_kernel="A_kernel"
            part_dtb="A_kernel-dtb"
        fi

        install_kernel "${part_kernel}"
        install_kernel "${part_dtb}"
    else
        echo "ERROR. ${NV_PAYLOAD_UPDATER} is NOT executable."
        echo "Post install script FAILED."
        echo "Cannot install package. Exiting..."
        exit 1
    fi
}

kernel_update () {
    local rootfs_ab=""

    if [ -x "${NVBOOTCTRL}" ]; then
        if [ -f "${payload_dir}/${kernel_only_payload}" ]; then
            pushd "${payload_dir}" &> /dev/null || exit 1

            # The payload updater currently only consumes payloads
            # named "/opt/ota_package/bl_update_payload"
            cp "${kernel_only_payload}" \
                "${OTA_PACKAGE_DIR}"/bl_update_payload

            popd &> /dev/null || exit 1
        else
            echo "ERROR. ${payload_dir}/${kernel_only_payload} does not exist."
            echo "Cannot install package. Exiting..."
            exit 1
        fi

        rootfs_ab=$( "${NVBOOTCTRL}" -t rootfs get-number-slots )
        if [ "${rootfs_ab}" = "1" ]; then
            echo "Rootfs AB is not enabled."
            # Update current kernel slot A and kernel-dtb slot A.
            install_kernel_on_rootfs_id "0"

        elif [ "${rootfs_ab}" = "2" ]; then
            echo "ERROR. Rootfs AB is enabled. Should never reach here."
            exit 1
        else
            echo "ERROR. Failed to get number of rootfs slots. Exiting..."
            exit 1
        fi

        echo "Reboot the target system for changes to take effect."
    else
        echo "ERROR. ${NVBOOTCTRL} is NOT executable."
        echo "Post install script FAILED."
        echo "Cannot install package. Exiting..."
        exit 1
    fi
}

is_overlayfs_enabled() {
	local efivar_dir="/sys/firmware/efi/efivars"
	local varname="L4TOverlayFsMode-360a04b9-3fe6-42c8-9fad-d9bf459c4770"

	# if no varialbe is defined it means overlays was not enabled
	if [ ! -f "${efivar_dir}/${varname}" ]; then
		OVERLAYFS_IS_ENABLED=0
	else
		val=$(xxd -ps "${efivar_dir}/${varname}")
		if [ "${val}" = "0700000001000000" ]; then
			OVERLAYFS_IS_ENABLED=1
		else
			OVERLAYFS_IS_ENABLED=0
		fi
	fi
}

if [ ! -r "${BOOT_CTRL_CONF}" ]; then
    echo "ERROR. Cannot open ${BOOT_CTRL_CONF} for reading."
    echo "Cannot install package. Exiting..."
    exit 1
fi

is_overlayfs_enabled
if [ "${OVERLAYFS_IS_ENABLED}" -eq 1 ]; then
	echo "Overlayfs is enabled. Ignore kernel partition update."
	exit 0
fi

"${UPDATE_CONF}" -l

chipid=$( awk '/TEGRA_CHIPID/ {print $2}' "${BOOT_CTRL_CONF}" )
payload_dir="${OTA_PACKAGE_DIR}"

case "${chipid}" in
    0x23)
        payload_dir+="/t23x"
        ;;
    0x26)
        payload_dir+="/t26x"
        echo "Skip update kernel partitions for T264."
        exit 0
        ;;
    *)
        echo "ERROR. Unrecognized chip ID: ${chipid}."
        echo "Cannot install kernel package. Exiting..."
        exit 1
esac

# Compare the compatible spec and tnspec to select the kernel payload
compare_specs_to_select_payloads

# Support using capsule update bootloader
echo "Starting kernel post-install procedure."
kernel_update

exit 0