#!/bin/bash
# mk-rpi4-sd.sh - Build a bootable SD card image for the Raspberry Pi 4/400
#
# Usage: ./scripts/mk-rpi4-sd.sh <kernel> <ramdisk.tar> <output.img> [ext2|ramdisk] [overlay]
#        (normally via `make ARCH=aarch64 rpi4-sd [ROOTFS=ramdisk]`)
#
# Root filesystem (4th argument, default ext2):
#   ext2     The system tree (the contents of ramdisk.tar) goes on a second,
#            ext2 partition labelled MONTAUKROOT, which the kernel mounts as
#            0:/ when Limine hands it no ramdisk (Fs/Boot.cpp). Limine only
#            loads the kernel, so boot no longer waits on the firmware's SD
#            driver reading the whole tree; files load on demand through
#            Drivers/Storage/Sdhci. Changes to 0:/ persist across boots.
#            ROOT_FREE_MIB (default 64) sets the free space on it; the
#            image is sized to its contents, not to a card. On first boot
#            0:/os/firstboot.elf grows the partition and the volume to fill
#            the card (it comes from the overlay directory, see below).
#   ramdisk  The old layout: Limine loads boot/ramdisk.tar from the ESP. Slow
#            to boot, but needs no SD driver in the kernel.
#
# Overlay (5th argument, ext2 only): a directory laid over the system tree,
# holding the files only the SD image ships (programs/bin-aarch64-rpi:
# os/firstboot.elf). Its first-boot service and capability grant are added to
# init.toml and capabilities.toml here, so the generic configs never name it.
#
# Boot chain: Pi EEPROM -> the Pi Firmware Task Force UEFI firmware
# (github.com/pftf/RPi4, an EDK2 build) on the card -> Limine
# (EFI/BOOT/BOOTAA64.EFI) -> the kernel. The firmware is downloaded once
# into rpi4-firmware/<version>/ (override with PFTF_VERSION).
#
# Partition 1 is MBR type 0xEF formatted FAT32, which is what the firmware
# boots from; with the ext2 root, partition 2 (type 0x83) follows it. Write
# the image with e.g.
#   sudo dd if=montauk-rpi4.img of=/dev/sdX bs=4M conv=fsync
#
# Firmware settings (UEFI setup, Esc at boot) MontaukOS relies on are the
# defaults: System Table Mode "ACPI" and "Limit RAM to 3 GB" enabled (the
# BCM2711's PCIe can only DMA below 3 GB, see the _DMA objects in its ACPI
# tables). The kernel console is mirrored to the PL011 on GPIO 14/15
# (pins 8/10, 115200 8N1; config.txt's miniuart-bt overlay moves the PL011
# there).

set -euo pipefail

if [ $# -lt 3 ] || [ $# -gt 5 ]; then
    echo "usage: $0 <kernel> <ramdisk.tar> <output.img> [ext2|ramdisk] [overlay]" >&2
    exit 1
fi

KERNEL="$1"
RAMDISK="$2"
OUTPUT="$3"
ROOTFS="${4:-ext2}"
OVERLAY="${5:-}"
ROOT_FREE_MIB="${ROOT_FREE_MIB:-64}"
ROOT_LABEL=MONTAUKROOT   # must match RootLabel in kernel/src/Fs/Boot.cpp

case "$ROOTFS" in
    ext2|ramdisk) ;;
    *) echo "mk-rpi4-sd: unknown root filesystem '$ROOTFS' (ext2 or ramdisk)" >&2; exit 1 ;;
esac

ROOT="$(cd "$(dirname "$0")/.." && pwd)"
PFTF_VERSION="${PFTF_VERSION:-v1.53}"
FW_DIR="$ROOT/rpi4-firmware/$PFTF_VERSION"
FW_ZIP="RPi4_UEFI_Firmware_${PFTF_VERSION}.zip"
FW_URL="https://github.com/pftf/RPi4/releases/download/${PFTF_VERSION}/${FW_ZIP}"

if [ ! -f "$FW_DIR/RPI_EFI.fd" ]; then
    echo "mk-rpi4-sd: downloading $FW_ZIP"
    mkdir -p "$FW_DIR"
    curl -fL --progress-bar -o "$FW_DIR/$FW_ZIP" "$FW_URL"
    unzip -q -o "$FW_DIR/$FW_ZIP" -d "$FW_DIR"
    rm -f "$FW_DIR/$FW_ZIP"
fi

WORK="$(mktemp -d)"
trap 'rm -rf "$WORK"' EXIT

# Register os/firstboot.elf (from the overlay) in a staged tree: an init service
# ahead of every other one, optional so init skips it quietly once the program
# has deleted itself, and the capabilities it needs to grow the volume
# (storage_admin) and delete itself from 0:/os (system_image).
add_firstboot_config() {
    local root="$1"
    [ -f "$root/os/firstboot.elf" ] || return 0
    local init="$root/config/init.toml" caps="$root/config/capabilities.toml"
    if ! grep -q '^\[services\.' "$init"; then
        echo "mk-rpi4-sd: no [services.*] table in $init" >&2
        exit 1
    fi
    awk '!done && /^\[services\./ {
            print "# Raspberry Pi SD image: grow the system volume to fill the card."
            print "# firstboot.elf deletes itself once that is done; this entry is then"
            print "# skipped."
            print "[services.firstboot]"
            print "path = \"0:/os/firstboot.elf\""
            print "name = \"first boot setup\""
            print "optional = true"
            print ""
            done = 1
        }
        { print }' "$init" > "$WORK/init.toml"
    cat "$WORK/init.toml" > "$init"
    cat >> "$caps" <<'TOML'

# ==== Raspberry Pi SD image ====

# Grows the system volume on first boot, then deletes itself.
[grant.firstboot]
path = "0:/os/firstboot.elf"
effective = ["storage_admin", "system_image"]
TOML
}

# ESP at 1 MiB: its payload plus room for a few kernel updates, at least
# 40 MiB so mformat can make FAT32 (65525 clusters). It holds the ramdisk
# only in ramdisk mode.
esp_payload=$(( $(stat -c %s "$KERNEL") + $(du -sb "$FW_DIR" | cut -f1) ))
[ "$ROOTFS" = ramdisk ] && esp_payload=$(( esp_payload + $(stat -c %s "$RAMDISK") ))
esp_mib=$(( esp_payload / 1048576 + 16 ))
[ $esp_mib -lt 40 ] && esp_mib=40

root_mib=0
if [ "$ROOTFS" = ext2 ]; then
    # Build the root from ramdisk.tar rather than the staged tree so both modes
    # ship identical contents (mkramdisk.sh has already resolved symlinks).
    mkdir "$WORK/root"
    tar -xpf "$RAMDISK" -C "$WORK/root"
    if [ -n "$OVERLAY" ] && [ -d "$OVERLAY" ]; then
        cp -rL "$OVERLAY"/. "$WORK/root"/
        add_firstboot_config "$WORK/root"
    else
        echo "mk-rpi4-sd: WARNING: no overlay${OVERLAY:+ at $OVERLAY}; the root will not grow to fill the card" >&2
    fi
    tree_kib=$(du -sk --apparent-size "$WORK/root" | cut -f1)
    # 4 KiB blocks round every file up and the inode tables take ~1.6%:
    # allow 15% plus 8 MiB for that, then the requested free space.
    root_mib=$(( tree_kib * 115 / 100 / 1024 + 8 + ROOT_FREE_MIB ))

    # Only features Fs/Ext2.cpp handles: it refuses any incompat flag but
    # filetype, and its buffers cap the block size at 4096. resize_inode
    # reserves room for the descriptor table to grow to 1 TiB (2^28 blocks),
    # which the first-boot grow needs (Fs/Ext2.cpp, "Online grow"). One inode
    # per 256 KiB keeps the inode tables the grow has to zero small: 0.1% of
    # the card (128 MB on 128 GB), still ~120k inodes on 32 GB.
    PATH=$PATH:/usr/sbin:/sbin mke2fs -q -F -t ext2 \
        -O none,filetype,sparse_super,resize_inode -E resize=268435456,root_owner=0:0 \
        -b 4096 -i 262144 -m 0 -L "$ROOT_LABEL" \
        -d "$WORK/root" "$WORK/root.ext2" "${root_mib}M"
fi

rm -f "$OUTPUT"
dd if=/dev/zero of="$OUTPUT" bs=1M count=0 seek=$((1 + esp_mib + root_mib)) status=none
if [ "$ROOTFS" = ext2 ]; then
    printf 'label: dos\nstart=2048, size=%dMiB, type=ef\nstart=%dMiB, type=83\n' \
        "$esp_mib" "$((1 + esp_mib))"
else
    printf 'label: dos\nstart=2048, type=ef\n'
fi | PATH=$PATH:/usr/sbin:/sbin sfdisk --quiet "$OUTPUT"

# The ESP is built as its own file: mtools sizes a filesystem at an @@offset
# to the end of the image, which would run the FAT over the root partition.
IMG="$WORK/esp.fat"
dd if=/dev/zero of="$IMG" bs=1M count=0 seek=$esp_mib status=none
mformat -i "$IMG" -F -v MONTAUK ::

# Firmware: everything in the release except its readme.
for f in "$FW_DIR"/*; do
    case "$(basename "$f")" in
        Readme.md) ;;
        *) mcopy -s -i "$IMG" "$f" :: ;;
    esac
done

mmd -i "$IMG" ::/EFI ::/EFI/BOOT ::/boot ::/boot/limine
mcopy -i "$IMG" "$ROOT/limine/BOOTAA64.EFI" ::/EFI/BOOT/
mcopy -i "$IMG" "$KERNEL" ::/boot/kernel

if [ "$ROOTFS" = ext2 ]; then
    # limine.conf without the ramdisk module: its absence is what makes the
    # kernel look for the MONTAUKROOT partition.
    grep -v -e '^[[:space:]]*module_path:' -e '^[[:space:]]*module_string:' \
        -e '^[[:space:]]*# Ramdisk module' "$ROOT/limine.conf" > "$WORK/limine.conf"
    mcopy -i "$IMG" "$WORK/limine.conf" ::/boot/limine/
    dd if="$WORK/root.ext2" of="$OUTPUT" bs=1M seek=$((1 + esp_mib)) conv=notrunc,sparse status=none
else
    mcopy -i "$IMG" "$ROOT/limine.conf" ::/boot/limine/
    mcopy -i "$IMG" "$RAMDISK" ::/boot/ramdisk.tar
fi
dd if="$IMG" of="$OUTPUT" bs=1M seek=1 conv=notrunc,sparse status=none

echo "mk-rpi4-sd: created $OUTPUT ($((1 + esp_mib + root_mib)) MiB, $ROOTFS root, firmware $PFTF_VERSION)"