#!/usr/bin/env bash
# Build provisioning initramfs for Raspberry Pi 4 / CM4 (aarch64).
# Produces initrd.img (gzip cpio) for TFTP boot (config.txt: initramfs initrd.img followkernel).
#
# Uses an Alpine Linux aarch64 base with Python 3 so rpi-eeprom-config can run
# directly in the initramfs (no chroot into eMMC needed for rescue).
#
# Requires Docker or Podman with arm64 emulation:
#   Fedora: sudo dnf install -y qemu-user-static
#   Debian/Ubuntu: sudo apt install -y qemu-user-static

set -e
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
OUTPUT="${1:-$SCRIPT_DIR/initrd.img}"
ALPINE_VERSION="3.22"
RPI_EEPROM_REPO="https://github.com/raspberrypi/rpi-eeprom"
RPI_EEPROM_RAW="https://raw.githubusercontent.com/raspberrypi/rpi-eeprom"

# Find container runtime
CONTAINER_RUNTIME=""
for cmd in podman docker; do
  if command -v "$cmd" >/dev/null 2>&1; then
    CONTAINER_RUNTIME="$cmd"
    break
  fi
done
if [ -z "$CONTAINER_RUNTIME" ]; then
  echo "Error: Docker or Podman is required to build the Alpine-based initramfs."
  echo "Install one of them and ensure arm64 emulation is available:"
  echo "  Fedora: sudo dnf install -y podman qemu-user-static"
  echo "  Debian/Ubuntu: sudo apt install -y podman qemu-user-static"
  exit 1
fi

echo "Using $CONTAINER_RUNTIME to build Alpine aarch64 initramfs..."

# Revision shown on serial so you can confirm the device is running the latest initrd
REV=$(date +%Y%m%d-%H%M 2>/dev/null || echo "unknown")
if [ -d "$SCRIPT_DIR/../.git" ] || [ -d "$SCRIPT_DIR/../../.git" ]; then
  REV="${REV}-$(git -C "$SCRIPT_DIR" rev-parse --short HEAD 2>/dev/null || echo 'nogit')"
fi

CNT_NAME="cm4-initramfs-build-$$"
trap "$CONTAINER_RUNTIME rm -f $CNT_NAME >/dev/null 2>&1; true" EXIT

# Build the rootfs inside an arm64 Alpine container
$CONTAINER_RUNTIME run --name "$CNT_NAME" --platform linux/arm64 \
  "alpine:${ALPINE_VERSION}" /bin/sh -c "
set -e

# Install packages: Python for rpi-eeprom-config, curl for provisioning client, coreutils for dd
apk add --no-cache python3 curl busybox coreutils

# Download rpi-eeprom tools from GitHub (use curl -sL which follows redirects; busybox wget does not)
curl -sL -o /usr/bin/rpi-eeprom-config \
  '${RPI_EEPROM_RAW}/master/rpi-eeprom-config'
curl -sL -o /usr/bin/rpi-eeprom-update \
  '${RPI_EEPROM_RAW}/master/rpi-eeprom-update'
curl -sL -o /usr/bin/rpi-eeprom-digest \
  '${RPI_EEPROM_RAW}/master/rpi-eeprom-digest'
chmod +x /usr/bin/rpi-eeprom-config /usr/bin/rpi-eeprom-update /usr/bin/rpi-eeprom-digest

# Download latest stable CM4 (BCM2711) EEPROM firmware
mkdir -p /lib/firmware/raspberrypi/bootloader/default
# Use GitHub API to find the latest pieeprom-*.bin in firmware-2711/default/
LATEST_FW=\$(curl -sL 'https://api.github.com/repos/raspberrypi/rpi-eeprom/contents/firmware-2711/default' \
  | grep -o '\"name\" *: *\"pieeprom-[^\"]*\\.bin\"' | sed 's/\"name\" *: *\"//;s/\"//' | sort | tail -1)
if [ -n \"\$LATEST_FW\" ]; then
  echo \"Downloading EEPROM firmware: \$LATEST_FW\"
  curl -sL -o /lib/firmware/raspberrypi/bootloader/default/pieeprom.bin \
    \"${RPI_EEPROM_RAW}/master/firmware-2711/default/\$LATEST_FW\"
else
  echo 'WARNING: Could not determine latest firmware; rescue-eeprom.sh may not work'
fi

# Create required directories
mkdir -p /proc /sys /dev/pts /mnt /tmp /run /usr/share/udhcpc /etc/default

# Clean up unnecessary files to reduce size
rm -rf /var/cache/apk/* /usr/share/man /usr/share/doc /usr/include \
  /usr/lib/python*/test /usr/lib/python*/unittest /usr/lib/python*/idlelib \
  /usr/lib/python*/tkinter /usr/lib/python*/ensurepip /usr/lib/python*/__pycache__/test* \
  /usr/lib/python*/lib2to3 /usr/lib/python*/distutils \
  /usr/share/terminfo/[a-s]* /usr/share/terminfo/[u-z]*

echo 'Alpine rootfs ready'
"

echo "Container build done. Copying scripts into container..."

# Copy our scripts into the container
$CONTAINER_RUNTIME cp "$SCRIPT_DIR/init" "$CNT_NAME:/init"
$CONTAINER_RUNTIME cp "$SCRIPT_DIR/provisioning-client.sh" "$CNT_NAME:/provisioning-client.sh"
$CONTAINER_RUNTIME cp "$SCRIPT_DIR/rescue-eeprom.sh" "$CNT_NAME:/rescue-eeprom.sh"
$CONTAINER_RUNTIME cp "$SCRIPT_DIR/udhcpc.script" "$CNT_NAME:/usr/share/udhcpc/default.script"

# Write revision.txt
echo "$REV" | $CONTAINER_RUNTIME cp /dev/stdin "$CNT_NAME:/revision.txt" 2>/dev/null || \
  (echo "$REV" > /tmp/cm4-rev-$$.txt && $CONTAINER_RUNTIME cp /tmp/cm4-rev-$$.txt "$CNT_NAME:/revision.txt" && rm -f /tmp/cm4-rev-$$.txt)

# Set permissions
$CONTAINER_RUNTIME start "$CNT_NAME" >/dev/null 2>&1 || true
$CONTAINER_RUNTIME exec "$CNT_NAME" chmod +x /init /provisioning-client.sh /rescue-eeprom.sh /usr/share/udhcpc/default.script 2>/dev/null || \
  echo "Note: could not chmod in stopped container; permissions set by cp"

# Export container filesystem and create cpio archive
echo "Exporting filesystem and building initrd.img..."
$CONTAINER_RUNTIME export "$CNT_NAME" | gzip -1 > /tmp/cm4-rootfs-$$.tar.gz

BUILD_DIR=$(mktemp -d)
trap "$CONTAINER_RUNTIME rm -f $CNT_NAME >/dev/null 2>&1; rm -rf $BUILD_DIR /tmp/cm4-rootfs-$$.tar.gz; true" EXIT

cd "$BUILD_DIR"
tar xzf /tmp/cm4-rootfs-$$.tar.gz
rm -f /tmp/cm4-rootfs-$$.tar.gz

# Remove container metadata that shouldn't be in initramfs
rm -rf .dockerenv .containerenv

# Ensure init is executable and at the root
chmod +x init provisioning-client.sh rescue-eeprom.sh usr/share/udhcpc/default.script 2>/dev/null || true

# Build cpio (gzip)
echo "Building cpio..."
find . -print0 | cpio -o -H newc -0 2>/dev/null | gzip -9 > "$OUTPUT"

SIZE=$(stat -c%s "$OUTPUT" 2>/dev/null || stat -f%z "$OUTPUT" 2>/dev/null || echo "?")
echo ""
echo "Written: $OUTPUT ($SIZE bytes, $(( ${SIZE:-0} / 1048576 )) MB)"
echo ""
echo "Next: copy initrd.img to TFTP root (e.g. /srv/tftpboot on LXC) and in config.txt add:"
echo "  initramfs initrd.img followkernel"
echo "Default provisioning server: http://10.20.50.1:5000 (override with kernel cmdline: provisioning_server=http://...)"