#!/bin/bash
# install-jetson-agx-orin-image-from-root-shell.sh
# Purpose: Downloads JetPack 6.1 (L4T R36.4.0) components and generates a bootable disk image.
# This script should be run from a root shell on a Jetson or ARM64 Linux system.

set -e

# Configuration
L4T_VERSION="R36.4.0"
BSP_URL="https://developer.nvidia.com/downloads/embedded/l4t/r36_release_v4.0/release/Jetson_Linux_R36.4.0_aarch64.tbz2"
ROOTFS_URL="https://developer.nvidia.com/downloads/embedded/l4t/r36_release_v4.0/release/Tegra_Linux_Sample-Root-Filesystem_R36.4.0_aarch64.tbz2"
OUTPUT_IMAGE="jetpack6_agx_orin.img"
BOARD="jetson-agx-orin-devkit"

echo "=== Jetson AGX Orin Image Preparation (JetPack 6.1 / L4T $L4T_VERSION) ==="

# 1. Check for free space (~20GB recommended)
FREE_SPACE=$(df -m . | awk 'NR==2 {print $4}')
if [ "$FREE_SPACE" -lt 20000 ]; then
    echo "Error: Less than 20GB of free space detected. Please run this in a larger partition."
    exit 1
fi

# 2. Install dependencies
echo "--- Installing required tools ---"
apt-get update
apt-get install -y wget lbzip2 qemu-user-static sudo coreutils

# 3. Download components
echo "--- Downloading Driver Package (BSP) ---"
wget -c $BSP_URL -O Jetson_Linux_aarch64.tbz2

echo "--- Downloading Sample Root Filesystem ---"
wget -c $ROOTFS_URL -O Tegra_Linux_Sample-Root-Filesystem_aarch64.tbz2

# 4. Extract and Assemble
echo "--- Extracting BSP ---"
tar xf Jetson_Linux_aarch64.tbz2

echo "--- Extracting RootFS (this takes a while) ---"
cd Linux_for_Tegra/rootfs/
sudo tar xf ../../Tegra_Linux_Sample-Root-Filesystem_aarch64.tbz2
cd ..

echo "--- Applying NVIDIA Binaries ---"
sudo ./apply_binaries.sh

# 5. Create the Disk Image
echo "--- Generating Disk Image: $OUTPUT_IMAGE ---"
# -o: output file
# -b: board type
# -d: device (USB/SD)
sudo ./tools/jetson-disk-image-creator.sh -o ../$OUTPUT_IMAGE -b $BOARD -d USB

cd ..

echo "==========================================================="
echo " SUCCESS: Image created at $(pwd)/$OUTPUT_IMAGE"
echo "==========================================================="
echo "Next Steps:"
echo "1. Identify your target disk (e.g., /dev/sda for USB or /dev/nvme0n1):"
echo "   lsblk"
echo ""
echo "2. Flash the image (Replace /dev/sdX with your target device):"
echo "   dd if=$OUTPUT_IMAGE of=/dev/sdX bs=4M status=progress conv=fsync"
echo ""
echo "3. Insert the media into the Jetson and boot while pressing ESC to select the boot device in UEFI."
echo "==========================================================="