#!/usr/bin/env bash
##############################################################################
# SPDX-License-Identifier: BSD-3-Clause
# Copyright (C) 2025 Engineering Design Team, Inc.
# https://edt.com
#
# Download and prepare NVIDIA's Jetson Linux to be flashed to an EDT carrier
# board.
#
# [1] NVIDIA Jetson Linux 36.4.0
#     https://developer.nvidia.com/embedded/jetson-linux-r3640
#
##############################################################################

set -eo pipefail
set -u
set -E

SCRIPT_FNAME=$(basename $0)
SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )

function fn_indent {
    sed 's/^/  /';
}

function fn_failure {
  local lineno=$1
  echo "[$SCRIPT_FNAME] Failed at line $lineno"
}

function fn_confirm_or_die {
    select yn in "Yes" "No"; do
        case $yn in
            Yes ) echo "User confirmed. Continuing..."; break;;
            No ) exit 255;;
            *) echo "That is not a valid choice"; continue;;
        esac
    done
}

function fn_main {
    if [ $EUID -eq 0 ]; then
        echo "Please do NOT run as sudo or root. Sudo will be used only for necessary steps."
        exit 1
    fi

    echo -n "Check if this Host PC is running the recommended Ubuntu version... "
    ubuntu_codename=$(grep 'UBUNTU_CODENAME' /etc/os-release)
    if [[ "$ubuntu_codename" == "UBUNTU_CODENAME=focal" ]]; then
        echo "YES (focal)"
    elif [[ "$ubuntu_codename" == "UBUNTU_CODENAME=jammy" ]]; then
        echo "YES (jammy)"
    else
        echo "NO ($ubuntu_codename)"
        echo "[WARNING] One of the recommended operating systems was not detected."
        echo "Are you sure you want to continue?"
        fn_confirm_or_die
    fi

    echo "Install system dependencies... "
    #sudo apt-get update | fn_indent
    sudo apt-get install curl lbzip2 sed | fn_indent

    # This is the "Driver Package (BSP)" link in the Jetson Linux
    # "Downloads and Links" list [1].
    echo -n "Download Jetson Linux (L4T)... "
    l4t_fname="Jetson_Linux_R36.4.4_aarch64.tbz2"
    if [ ! -f "$l4t_fname" ]; then
        echo ""
        curl --progress-bar \
            -o "$l4t_fname" \
            --location \
            'https://developer.nvidia.com/downloads/embedded/l4t/r36_release_v4.4/release/Jetson_Linux_r36.4.4_aarch64.tbz2';
    else
        echo "SKIPPED"
    fi
    echo "Checking L4T archive integrity..."
    echo "a6ce11c22100ab0976e419959182417c83d62f4272501bc8714f2e076e010f3b  ./${l4t_fname}" > checksum.sha256
    sha256sum --check ./checksum.sha256 | fn_indent
    rm checksum.sha256

    # This is the "Sample Root Filesystem" link in the Jetson Linux "Downloads
    # and Links" list [1]. The sample root filesystem contains the system
    # libraries and applications that make up the Ubuntu operating system.
    echo -n "Download Sample Root Filesystem (rootfs)... "
    rootfs_fname="Tegra_Linux_Sample-Root-Filesystem_R36.4.4_aarch64.tbz2"
    if [ ! -f "$rootfs_fname" ]; then
        echo ""
        curl --progress-bar \
            -o "$rootfs_fname" \
            --location \
            'https://developer.nvidia.com/downloads/embedded/l4t/r36_release_v4.4/release/Tegra_Linux_Sample-Root-Filesystem_r36.4.4_aarch64.tbz2';
    else
        echo "SKIPPED"
    fi

    echo "Checking rootfs archive integrity... "
    echo "2983173288a14a9eb3e82531b30b61009ea9cf5634ac4a16f896df8208a7359b  ./${rootfs_fname}" > checksum.sha256
    sha256sum --check ./checksum.sha256 | fn_indent
    rm checksum.sha256

    L4T_DIR="$PWD/Linux_for_Tegra"

    echo -n "Unpacking Jetson Linux (L4T)... "
    if [ ! -f "$L4T_DIR/apply_binaries.sh" ]; then
        echo ""
        # Use lbzip2 to speed this up
        tar -x --checkpoint=.100 --use-compress-program=lbzip2 -f "$l4t_fname"
        echo ""
    else
        echo "SKIPPED"
    fi

    # Extract the sample RootFS inside of the "Linux_for_Tegra/" directory created
    # in the previous step. These files should have root ownership, so extract with
    # sudo.
    echo -n "Unpacking rootfs into Linux_for_Tegra/ directory... "
    if [ ! -f "$L4T_DIR/rootfs/etc/hosts" ]; then
        echo ""
        echo "(sudo) Unpack sample rootfs..."
        sudo tar -x --checkpoint=.1000 --use-compress-program=lbzip2 \
            -C "$L4T_DIR/rootfs" \
            -f "./$rootfs_fname"
        echo ""
    else
        echo "SKIPPED"
    fi

    sudo rm -fv ${L4T_DIR}/rootfs/README*

    # Ensure your Host PC has NVIDIA's flashing pre-requisites.
    echo "Run 'l4t_flash_prerequisites.sh' in L4T... "
    sudo "$L4T_DIR"/tools/l4t_flash_prerequisites.sh | fn_indent

    # Copy the board-specific customizations from this package's "l4t/" directory
    # to the "Linux_for_Tegra/" directory.
    echo "Copy EDT's customized L4T files into Linux_for_Tegra/... "
    cp --preserve=mode,timestamps --recursive  --force --verbose "$SCRIPT_DIR"/l4t/* "$L4T_DIR" | fn_indent

    pushd "$L4T_DIR"

    # Apply system packages for the target to the staged root file system.
    echo "Run 'apply_binaries.sh' in L4T... "
    sudo ./apply_binaries.sh | fn_indent

    # Copy the board-specific customizations from this package's "rootfs/" directory
    # to the "Linux_for_Tegra/rootfs" directory.
    echo "Copy EDT's customized RootFS files into Linux_for_Tegra/rootfs/... "
    sudo cp --preserve=mode,timestamps --recursive --force --verbose "$SCRIPT_DIR"/rootfs "$L4T_DIR" | fn_indent

    fn_board_specific_final_changes

    # Run NVIDIA's flash script with the "--no-flash" argument to apply the
    # remaining customizations and build system image binaries.

    if [ -z ${BOARDREV+x} ]; then
        BOARDREV="0.0"
    fi

    echo ""
    echo "Building system images for flashing, by running:"
    echo "    sudo BOARDID=${BOARDID} BOARDSKU=${BOARDSKU} BOARDREV=${BOARDREV} FAB=${FAB} CHIP_SKU="${CHIP_SKU}" RAMCODE_ID=${RAMCODE_ID} \ "
    echo "        ${EDT_BOARD_FLASH_SH} --no-flash " "${EDT_BOARD_FLASH_SH_ARGS[@]}"
    echo ""
    sudo BOARDID="${BOARDID}" BOARDSKU="${BOARDSKU}" BOARDREV="${BOARDREV}" FAB="${FAB}" CHIP_SKU="${CHIP_SKU}" RAMCODE_ID="${RAMCODE_ID}" \
        "$EDT_BOARD_FLASH_SH" --no-flash "${EDT_BOARD_FLASH_SH_ARGS[@]}" | fn_indent

    echo "[$SCRIPT_FNAME] Done!"

    echo "[$SCRIPT_FNAME] Next steps:"
    echo "    - Put the target device into Recovery Mode."
    echo "    - Connect this Host PC to the target device's specific connector used for Recovery flashing."
    echo "    - Run ./edt_bsp_flash.sh in this directory."
}

trap 'fn_failure ${LINENO}' ERR

# Include board-specific variables and functions.
source "$SCRIPT_DIR"/_bsp_script_vars.sh

fn_main | tee "$SCRIPT_FNAME".log
