#!/bin/bash

# SPDX-FileCopyrightText: Copyright (c) 2020-2023 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"
T186REF_UPDATER="/usr/sbin/nv_update_engine"
T186REF_NVBOOTCTRL="/usr/sbin/nvbootctrl"
UPDATE_CONF="/opt/nvidia/l4t-bootloader-config/nv-l4t-bootloader-config.sh"
L4T_LAUNCHER="BOOTAA64.efi"
EMMC_VAR_MOUNT_DIR="/opt/nvidia/emmc_var"
ACTIVE_ESP_MOUNT_DIR="/opt/nvidia/esp"
ESP_CAPSULE_DIR="EFI/UpdateCapsule/"
ESP_LAUNCHER_DIR="EFI/BOOT/"
ESP_VAR_DIR="EFI/NVDA/Variables"
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"
bl_only_payload=""
capsule_payload=""

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

chkerr () {
    # shellcheck disable=SC2181
    if [ $? -ne 0 ]; then
        if [ "$1" != "" ]; then
            echo "$1";
        else
            echo "failed.";
        fi;
        exit 1;
    fi;
    if [ "$1" = "" ]; then
        echo "done.";
    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-)
            bl_only_payload="bl_only_payload_3701_agx_ind"
            capsule_payload="TEGRA_BL_3701_agx_ind.Cap"
            ;;
        *jetson-agx-orin-devkit- | *jetson-agx-orin-devkit-maxn-)
            bl_only_payload="bl_only_payload_3701_agx"
            capsule_payload="TEGRA_BL_3701_agx.Cap"
            ;;
        *holoscan-devkit-)
            bl_only_payload="bl_only_payload_3701_holoscan"
            capsule_payload="TEGRA_BL_3701_holoscan.Cap"
            ;;
        *igx-orin-devkit-)
            bl_only_payload="bl_only_payload_3701_igx"
            capsule_payload="TEGRA_BL_3701_igx.Cap"
            ;;
        *pjai*)
            bl_only_payload="bl_only_payload"
            capsule_payload="TEGRA_BL.Cap"
            ;;
        *)
            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
    bl_only_payload="bl_only_payload_3701"
    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."
        bl_only_payload="bl_only_payload_3701_000"
        capsule_payload="TEGRA_BL_3701_000.Cap"
    else
        echo "INFO. ${board_fab} does not match any known boards."
    fi
}

select_3767_payload () {
    bl_only_payload="bl_only_payload"
    capsule_payload="TEGRA_BL.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
    bl_only_payload="bl_only_payload"
    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
}

# Mount esp on specified media to specified directory
mount_media_esp () {
    local media_dev="${1}"
    local mount_dir="${2}"
    local esp_uuid=""

    # Create esp mount directory if does not exist
    if [ ! -d "${mount_dir}" ]; then
        mkdir -p "${mount_dir}"
    fi

    # Umount the esp mount directory in case it is mounted
    if mountpoint -q "${mount_dir}"; then
        umount "${mount_dir}" > /dev/null 2>&1
    fi

    # Find the uuid of specified media's esp partition
    esp_uuid=$( lsblk -o partlabel,uuid,name | grep "${media_dev}" | grep -w "esp" | awk '{ print $2 }' )
    if [ "${esp_uuid}" == "" ]; then
        echo "Error: uuid of esp is not found."
        return 1
    fi

    # Unmount the esp in case it is mounted to other directory
    umount /dev/disk/by-uuid/"${esp_uuid}" > /dev/null 2>&1

    # Mount the esp to specified ${mount_dir}
    if ! mount UUID="${esp_uuid}" "${mount_dir}"; then
        return 2
    fi

    return 0
}

extlinux_update () {
	local extlinux_conf="/boot/extlinux/extlinux.conf"

	echo "Updating extlinux.conf..."
	# check whether exlinux.conf exists
	if [ ! -f "${extlinux_conf}" ]; then
		if [ "${chipid}" == "0x19" ]; then
			echo "ERROR. ${extlinux_conf} does not exist"
			exit 1
		fi
		echo "File ${extlinux_conf} does not exist"
		return 0
	fi

	# check whether root= is set in extlinut.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"
		return 0
	fi

	# 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
}

force_install_payload () {
    local payload_name="${1}"

    echo "INFO. Force install payload: ${payload_name}."
    if [ -f "${payload_dir}/${payload_name}" ]; then
        # shellcheck disable=SC2164
        pushd "${payload_dir}" &> /dev/null
        # The update engine and payload updater currently only consumes payloads
        # named "/opt/ota_package/bl_update_payload"
        cp "${payload_name}" "${OTA_PACKAGE_DIR}"/bl_update_payload
        "${T186REF_UPDATER}" --forced-install &> "${OTA_PACKAGE_DIR}/${payload_name}.log"
        # shellcheck disable=SC2181
        if [ $? -ne 0 ]; then
            echo "ERROR. Procedure for ${payload_name} update FAILED."
            echo "Cannot install package. Exiting..."
            exit 1
        fi
        # shellcheck disable=SC2164
        popd &> /dev/null
        echo "Update ${payload_name} completed."
    else
        echo "ERROR. ${payload_dir}/${payload_name} does not exist."
        echo "Cannot install package. Exiting..."
        exit 1
    fi
}

force_install_esp() {
    local payload_name="${1}";
    local l4tlauncher_dir=""

    if [ -f "${payload_dir}/${payload_name}" ]; then
        l4tlauncher_dir="${ACTIVE_ESP_MOUNT_DIR}/${ESP_LAUNCHER_DIR}"
        echo "Updating esp partition...";
        cp -f "${payload_dir}/${payload_name}" "${l4tlauncher_dir}";
        chkerr "Copying ${payload_dir}/${payload_name} failed.";
        sync;
        chkerr;
    else
        echo "ERROR. ${payload_dir}/${payload_name} does not exist.";
        echo "Cannot install package. Exiting...";
        exit 1;
    fi
}

force_update() {
    echo "Starting legacy update procedure."

    if [ -x "${T186REF_UPDATER}" ]; then
        # Install bootloader components to current slot
        # Installing bl_only_payload.
        # kernel and kernel-dtb are installed along with Kernel Deb package.
        force_install_payload "${bl_only_payload}"

        # Until UEFI capsule mechanism is ready, we should always update esp
        # partition in case L4tlauncher is changed. In general, esp partition update
        # is handled by UEFI.
        force_install_esp "${L4T_LAUNCHER}"

    else
        echo "ERROR. ${T186REF_UPDATER} or ${T186REF_NVBOOTCTRL} is NOT executable."
        echo "Post install script FAILED."
        echo "Cannot install package. Exiting..."
        exit 1
    fi
}

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

# Mount active boot media's esp to /opt/nvidia/esp directory
# Mount eMMC's esp to /opt/nvidia/emmc_var if board is eMMC only
# and the active boot media is not eMMC
mount_esp () {
    # Mount active boot media's esp to /opt/nvidia/esp
    if ! mount_media_esp "${active_media}" "${ACTIVE_ESP_MOUNT_DIR}"; then
        echo "ERROR. Mount active boot media's esp failed."
        echo "Cannot install capsule package. Exiting..."
        exit 1
    fi

    if [ "${emmc_only}" = "true" ] && [ "${active_media}" != "mmcblk0" ]; then
        # Mount eMMC's esp to /opt/nvidia/emmc_var directory
        if ! mount_media_esp "mmcblk0" "${EMMC_VAR_MOUNT_DIR}"; then
            echo "ERROR. Mount eMMC's esp partition failed."
            echo "Cannot install capsule package. Exiting..."
            exit 1
        fi
    fi

}

# Unmount the esp if mounted
unmount_esp () {
    # Unmount eMMC's esp
    if mountpoint -q "${EMMC_VAR_MOUNT_DIR}"; then
        umount "${EMMC_VAR_MOUNT_DIR}" > /dev/null 2>&1;
    fi

    # Unmount active boot media's esp
    if mountpoint -q "${ACTIVE_ESP_MOUNT_DIR}"; then
        umount "${ACTIVE_ESP_MOUNT_DIR}" > /dev/null 2>&1;
    fi
}

# 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
}

# 1. Copy the capsule payload to ${actvie_media's esp}/EFI/UpdateCapsule/;
# 2. Copy the L4Tlauncher to ${active_media's esp}/EFI/BOOT/.
copy_capsule_and_l4tlauncher () {
    local payload_name="${1}"
    local capsule_dir=""
    local l4tlauncher_dir=""

    echo "INFO. Copy capsule payload: ${payload_name}."
    # Check capsule payload, exit if does not exist
    if [ ! -f "${payload_dir}/${payload_name}" ]; then
        echo "ERROR. Cannot find ${payload_dir}/${payload_name}."
        echo "Cannot install package. Exiting..."
        exit 1
    fi

    capsule_dir="${ACTIVE_ESP_MOUNT_DIR}/${ESP_CAPSULE_DIR}"
    l4tlauncher_dir="${ACTIVE_ESP_MOUNT_DIR}/${ESP_LAUNCHER_DIR}"

    # Create UpdateCapsule directory if does not exist
    if [ ! -d "${capsule_dir}" ]; then
        mkdir -p "${capsule_dir}"
    fi

    # Copy capsule payload and L4Tlauncher
    # shellcheck disable=SC2164
    pushd "${payload_dir}" &> /dev/null
    cp "${payload_name}" "${capsule_dir}"
    cp "${L4T_LAUNCHER}" "${l4tlauncher_dir}"
    # shellcheck disable=SC2164
    popd &> /dev/null

    echo "INFO. Copy capsule payload to ${capsule_dir} and l4tlauncher to ${l4tlauncher_dir} done."
}

# Set bit 2 of the OsIndications UEFI variable
set_capsule_uefi_variable () {
    local var_path=""

    # Set the variable directory
    if [ "${emmc_only}" = "true" ]; then
        if [ "${active_media}" = "mmcblk0" ]; then
            var_path="${ACTIVE_ESP_MOUNT_DIR}/${ESP_VAR_DIR}/${OS_INDICATION_VAR}-${EFI_GLOBAL_GUID}"
        else
            var_path="${EMMC_VAR_MOUNT_DIR}/${ESP_VAR_DIR}/${OS_INDICATION_VAR}-${EFI_GLOBAL_GUID}"
        fi
    else
        var_path="${SYS_VAR_DIR}/${OS_INDICATION_VAR}-${EFI_GLOBAL_GUID}"
    fi

    if [ -f "${TMP_FILE}" ]; then
        rm -f "${TMP_FILE}"
    fi

    # Set bit 2 of the OsIndications UEFI variable
    printf "\x07\x00\x00\x00\x04\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. Set capsule UEFI variable ${var_path} done."
}

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

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

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

"${UPDATE_CONF}" -l

chipid=$( awk '/TEGRA_CHIPID/ {print $2}' "${BOOT_CTRL_CONF}" )
legacy_update_support=$( awk '/TEGRA_LEGACY_UPDATE/ {print $2}' "${BOOT_CTRL_CONF}" )
emmc_only=$( awk '/TEGRA_EMMC_ONLY/ {print $2}' "${BOOT_CTRL_CONF}" )
active_media=$( awk '/TEGRA_BOOT_STORAGE/ {print $2}' "${BOOT_CTRL_CONF}" )
payload_dir="${OTA_PACKAGE_DIR}"

set_payload_dir

mount_esp

# Dump debug info before update
dump_debug_info

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

# If legacy update is not supported, prepare capsule to update bootloader;
# if legacy update is supported, call nv_update_engine to update bootloader.
echo "Starting bootloader post-install procedure."
if [ "${legacy_update_support}" = "false" ]; then
    # Prepare capsule update
    copy_capsule_and_l4tlauncher ${capsule_payload}
    set_capsule_uefi_variable
elif [ "${legacy_update_support}" = "true" ]; then
    # Use BUP to update
    install_mtdblock
    force_update
else
    echo "ERROR. Unrecognized legacy update support: ${legacy_update_support}."
    echo "Cannot install bootloader package. Exiting..."
    exit 1
fi

unmount_esp

extlinux_update

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