#!/bin/bash

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

# This script runs on target post installation of nvidia-bootloader debian package.

BOOT_CTRL_CONF="/etc/nv_boot_control.conf"
OTA_PACKAGE_DIR="/opt/ota_package"
NVBOOTCTRL="/usr/sbin/nvbootctrl"
UPDATE_CONF="/opt/nvidia/l4t-bootloader-config/nv-l4t-bootloader-config.sh"
L4T_LAUNCHER="BOOTAA64.efi"
ACTIVE_ESP_MOUNT_DIR="/boot/efi"
ESP_LAUNCHER_DIR="EFI/BOOT/"
SYS_VAR_DIR="/sys/firmware/efi/efivars"
EFI_GLOBAL_GUID="8be4df61-93ca-11d2-aa0d-00e098032b8c"
TMP_FILE="/tmp/spec_var.bin"
OS_INDICATION_VAR="OsIndications"
capsule_payload=""
OVERLAYFS_IS_ENABLED=0
EXTLINUX_CONF="/boot/extlinux/extlinux.conf"

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

reboot_notify () {
    if [ -d /var/lib/update-notifier/user.d/ ]; then
        notifier_file=/var/lib/update-notifier/user.d/${DPKG_MAINTSCRIPT_PACKAGE}
        if [ -e "${notifier_file}" ];then
            rm -rf "${notifier_file}"
        fi
        cat << EOF >> "${notifier_file}"
Name: ${DPKG_MAINTSCRIPT_PACKAGE} Post Install Notification.
Priority: High
Terminal: False
DontShowAfterReboot: True
Description: Reboot is required to complete the installation.
EOF
        touch /var/lib/update-notifier/dpkg-run-stamp
    fi

    if [ -x /usr/share/update-notifier/notify-reboot-required ]; then
        /usr/share/update-notifier/notify-reboot-required
    fi
}

cmp_board_name_select_3701_payload () {
    local compat_spec="${1}"

    case "${compat_spec}" in
        *jetson-agx-orin-devkit-industrial- | *jetson-agx-orin-devkit-industrial-maxn-)
            capsule_payload="TEGRA_BL_3701_agx_ind.Cap"
            ;;
        *jetson-agx-orin-devkit- | *jetson-agx-orin-devkit-maxn-)
            capsule_payload="TEGRA_BL_3701_agx.Cap"
            ;;
        *holoscan-devkit-)
            capsule_payload="TEGRA_BL_3701_holoscan.Cap"
            ;;
        *igx-orin-devkit-)
            echo "INFO. Skip installing bootloader for igx-orin-devkit."
            exit 0
            ;;
        *)
            echo "ERROR. ${compat_spec} does not match any known boards."
            exit 1
            ;;
    esac
}

select_3701_payload () {
    local compat_spec="${1}"
    local tn_spec=""
    local board_fab=""
    local found="false"

    # Default payloads
    capsule_payload="TEGRA_BL_3701.Cap"

    # 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-'
    )
    if [ "${compat_spec}" = "" ]; then
        echo "ERROR. Invalid COMPATIBLE_SPEC."
        return 1
    fi

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

    # Curernt board is not in the zero fab list, select the payloads
    # according to the board name in the compatible spec.
    if [ "${found}" != "true" ]; then
        # Compare the compatible board name to select the payloads
        cmp_board_name_select_3701_payload "${compat_spec}"
        return 0
    fi

    # Current board is in the zero fab list, select the payloads
    # according to the board FAB.
    # Get the TNSPEC and parse board_fab in it
    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
    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 bootloader payload file."
        capsule_payload="TEGRA_BL_3701_000.Cap"
    else
        echo "INFO. ${board_fab} does not match any known boards."
    fi
}

select_3767_payload () {
    capsule_payload="TEGRA_BL_3767.Cap"
}

# Select BUP and Capsule payloads based on the compatible spec and tnspec
compare_specs_to_select_payloads () {
    local compatible_spec=""
    local board_id=""

    # Default payloads
    capsule_payload="TEGRA_BL.Cap"

    # BUP and Capsule payloads selection only apply for t23x
    if [ "${chipid}" != "0x23" ]; then
        return 0
    fi

    # Select the payloads according to the board ID.
    compatible_spec=$( awk '/COMPATIBLE_SPEC/ {print $2}' "${BOOT_CTRL_CONF}" )
    if [ "${compatible_spec}" = "" ]; then
        echo "ERROR. Cannot find COMPATIBLE_SPEC in ${BOOT_CTRL_CONF}."
        return 1
    fi
    board_id=$( echo "${compatible_spec}" | awk -F"-" '{print $1}' )
    case "${board_id}" in
        3701)
            select_3701_payload "${compatible_spec}"
            ;;
        3767)
            select_3767_payload
            ;;
        *)
            echo "ERROR. Unrecognized board ID: ${board_id}."
            echo "Exiting..."
            exit 1
    esac

    return 0
}

menu_line=""
linux_line=""
fdt_line=""
initrd_line=""
append_line=""
overlays_line=""

# Parse the primary entry of extlinux.conf
extlinux_parse () {
    # shellcheck disable=SC2162
    while read -a line
    do
        # Find the LABEL primary
        if [ ${#line[@]} -eq 2 ] && [ "${line[0]}" = "LABEL" ]; then
            label="${line[1]}"
        fi
        # Parse the primary entry
        if [ "${label}" = "primary" ]; then
            case "${line[0]}" in
                MENU) menu_line="${line[*]}";;
                LINUX) linux_line="${line[*]}";;
                FDT) fdt_line="${line[*]}";;
                INITRD) initrd_line="${line[*]}";;
                APPEND) append_line="${line[*]}";;
                OVERLAYS) overlays_line="${line[*]}";;
                *) ;;
            esac
        fi
    done <<<"$(cat "${EXTLINUX_CONF}")"
}

# Add a new primary-backup entry to the end of extlinux.conf
add_backup_entry ()
{
    echo "LABEL primary-backup" >> "${EXTLINUX_CONF}"
    if [ "${menu_line}" != "" ]; then
        echo "      ${menu_line} backup" >> "${EXTLINUX_CONF}"
    fi
    if [ "${linux_line}" != "" ]; then
        echo "      ${linux_line}" >> "${EXTLINUX_CONF}"
    fi
    if [ "${fdt_line}" != "" ]; then
        echo "      ${fdt_line}" >> "${EXTLINUX_CONF}"
    fi
    if [ "${initrd_line}" != "" ]; then
        echo "      ${initrd_line}" >> "${EXTLINUX_CONF}"
    fi
    if [ "${append_line}" != "" ]; then
        echo "      ${append_line}" >> "${EXTLINUX_CONF}"
    fi
    if [ "${overlays_line}" != "" ]; then
        echo "      ${overlays_line}" >> "${EXTLINUX_CONF}"
    fi
}

# Check whether FDT line is in the primary entry of extlinux.conf:
# If yes, backup the primary entry as primary-backup and remove the
# FDT line from the primary entry.
# If no, do nothing.
check_and_update_FDT () {
    # Parse the primary entry of extlinux.conf
    extlinux_parse

    if [ "${fdt_line}" != "" ]; then
        # Remove the FDT line from the primary entry
        sed -i '/.*FDT/d' "${EXTLINUX_CONF}"

        # Add a primary_backup entry at the end of extlinux.conf
        local primary_backup=""
        primary_backup=$(grep "^LABEL primary-backup" "${EXTLINUX_CONF}")
        if [ "${primary_backup}" != "" ]; then
            # Delete the existed primary_backup entry
            # Assume the primary_backup is always the last entry
            sed -i '/LABEL primary-backup/,$d' "${EXTLINUX_CONF}"
        fi
        add_backup_entry
        echo "Info: add backup entry done."
    fi
}

check_and_update_APPEND ()
{
    # check whether root= is set in extlinux.conf
    # shellcheck disable=SC2155
    local root_dev="$(grep -E "^[ \t]*APPEND.*[ \t]root=" ${EXTLINUX_CONF})"
    if [ "${root_dev}" != "" ]; then
        echo "Root device is set in the extlinux.conf"
    else
        # get bootargs from /proc/cmdline
        # shellcheck disable=SC2155
        local boot_args="$(sed "s/\(root=.*\)[ \t]video=.*/\1/" </proc/cmdline)"
        if [ "${boot_args}" == "" ]; then
            echo "Boot args is not found from /proc/cmdline"
            exit 1
        fi

        # fill the bootargs into the "APPEND" line in the extlinux.conf
        echo "Adding bootargs into exlinux.conf..."
        sed -i "/^[ \t]*APPEND/s|\$| ${boot_args}|" "${EXTLINUX_CONF}"
        root_dev="$(grep -E "^[ \t]*APPEND.*[ \t]root=" ${EXTLINUX_CONF})"
        if [ "${root_dev}" == "" ]; then
            echo "Failed to add boot args into the extlinux.conf"
            exit 1
        fi
    fi
}

extlinux_update () {
    echo "Updating extlinux.conf..."
    # check whether exlinux.conf exists
    if [ ! -f "${EXTLINUX_CONF}" ]; then
        echo "File ${EXTLINUX_CONF} does not exist"
        return 0
    fi

    # Check whether root= is set in APPEND line of extlinux.conf
    check_and_update_APPEND

    # Check whether FDT line is set in extlinux.conf
    check_and_update_FDT
}

# Set payload_dir according to the chipid
set_payload_dir () {
    case "${chipid}" in
        0x23)
            payload_dir+="/t23x"
            ;;
        *)
            echo "ERROR. Unrecognized chip ID: ${chipid}."
            echo "Cannot find payload directory. Exiting..."
            exit 1
    esac
}

# Install mtdblock and its dependent modules
install_mtdblock () {
    local kernel_ver=
    local mtd_dir=
    kernel_ver="$(uname -r)"
    mtd_dir="/lib/modules/${kernel_ver}/kernel/drivers/mtd"

    # mtdblock has dependency on mtd and mtd_blkdevs
    if [ -f "${mtd_dir}/mtd.ko" ] \
        &&  [ -f "${mtd_dir}/mtd_blkdevs.ko" ] \
        &&  [ -f "${mtd_dir}/mtdblock.ko" ]; then
        echo "INFO. Installing mtdblock."
        modprobe mtdblock
    fi
}

# Uninstall mtdblock and its dependent modules
uninstall_mtdblock () {
    echo "INFO. Uninstalling mtdblock."
    modprobe -r mtdblock
}

# Copy the L4Tlauncher to active boot media's esp /EFI/BOOT/ directory.
copy_l4tlauncher () {
    local l4tlauncher_dir=""

    l4tlauncher_dir="${ACTIVE_ESP_MOUNT_DIR}/${ESP_LAUNCHER_DIR}"
    # shellcheck disable=SC2164
    pushd "${payload_dir}" &> /dev/null
    cp "${L4T_LAUNCHER}" "${l4tlauncher_dir}"
    # shellcheck disable=SC2164
    popd &> /dev/null

    echo "INFO. Copy L4tlauncher to ${l4tlauncher_dir} done."
}

# Create the OsIndications UEFI variable if not exist
create_capsule_uefi_variable () {
    local var_path=""

    # Set the variable directory
    var_path="${SYS_VAR_DIR}/${OS_INDICATION_VAR}-${EFI_GLOBAL_GUID}"
    if [ -f "${var_path}" ]; then
        # Return if the variable exist
        return;
    fi

    # Create the OsIndications UEFI variable if it does not exist
    printf "\x07\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00" > "${TMP_FILE}"
    dd if="${TMP_FILE}" of="${var_path}" > /dev/null 2>&1

    # Clean up the tmp fie
    rm -f "${TMP_FILE}"

    echo "INFO. Create capsule UEFI variable ${var_path} done."
}

# Call the fwupdtool to trigger a capsule update
trigger_capsule_update () {
    # Check the utilities are installed
    if [ ! -f "$(command -v fwupdmgr)" ]; then
        echo "ERROR. fwupdmgr is not installed."
        exit 1
    fi
    if [ ! -f "$(command -v fwupdtool)" ]; then
        echo "ERROR. fwupdtool is not installed."
        exit 1
    fi

    # Since the nv-l4t-bootloader-config.sh is called at the beginning of
    # this script, it checks and mounts the active esp to /boot/efi,
    # so, here only check whether the /boot/efi is mount or not.
    if ! mountpoint -q "${ACTIVE_ESP_MOUNT_DIR}"; then
        echo "ERROR. The esp is not mounted to ${ACTIVE_ESP_MOUNT_DIR}."
        exit 1
    fi

    # Trigger the capsule update
    create_capsule_uefi_variable
    local device_id=""
    device_id=$(fwupdmgr get-devices | awk '/System Firmware/{getline; print}' | tr -cd '0-9 a-z A-Z\n' | awk '{print $3}')
    if [ "${device_id}" != "" ]; then
        # Call fwupdtool to copy capsule payload and set UEFI variable
        yes N | fwupdtool install-blob "${payload_dir}/${capsule_payload}" "${device_id}" > /dev/null 2>&1
    else
        echo "ERROR. NULL Device ID."
        exit 1
    fi

    echo "INFO. Trigger Capsule update is done."
}

dump_debug_info () {
    # Dump slots info
    if [ -x "${NVBOOTCTRL}" ]; then
        echo "INFO. Dump slots info:"
        "${NVBOOTCTRL}" dump-slots-info
    fi

    # Dump ${BOOT_CTRL_CONF}
    echo "INFO. Dump nv_boot_control.conf:"
    cat "${BOOT_CTRL_CONF}"
}

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

	# if no variable 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 bootloader partition update."
	exit 0
fi

"${UPDATE_CONF}" -l

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

set_payload_dir

# Dump debug info before update
dump_debug_info

# Compare the compatible spec and tnspec to select the bootloader payloads
compare_specs_to_select_payloads

echo "Starting bootloader post-install procedure."
# Trigger capsule update
trigger_capsule_update

# Copy the L4tlauncher to the ESP partition
copy_l4tlauncher

extlinux_update

echo "Reboot the target system for updates to take effect."
reboot_notify