#!/bin/bash
# flash-cm4.sh - Flash CM4 eMMC while EMMC_DISABLE jumper is bridged
# Usage: ./scripts/flash-cm4.sh [/dev/sdX]
#   If no device given, auto-detects the CM4 USB mass storage device.
set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
BEACON_DIR="$(cd "$SCRIPT_DIR/.." && pwd)"
REPO_ROOT="$(cd "$BEACON_DIR/.." && pwd)"
USBBOOT_DIR="$REPO_ROOT/usbboot"
MSD_DIR="$USBBOOT_DIR/secure-boot-msd"
IMAGE="$REPO_ROOT/output/images/sdcard.img.xz"
PRIVATE_KEY="$REPO_ROOT/private.pem"

# Build rpiboot from source if not already compiled
if [ ! -x "$USBBOOT_DIR/rpiboot" ]; then
    echo "==> Building rpiboot from source..."
    make -C "$USBBOOT_DIR"
fi

# Ensure secure-boot-msd boot.img is signed (required for secure-boot-locked CM4)
if [ ! -f "$MSD_DIR/boot.sig" ] || [ "$MSD_DIR/boot.img" -nt "$MSD_DIR/boot.sig" ]; then
    echo "==> Signing secure-boot-msd/boot.img with private.pem..."
    "$USBBOOT_DIR/tools/rpi-eeprom-digest" -i "$MSD_DIR/boot.img" \
        -o "$MSD_DIR/boot.sig" -k "$PRIVATE_KEY"
fi

find_removable_sd() {
    lsblk -dno NAME,RM | awk '$2==1{print $1}' | grep '^sd' | head -1
}

# Step 1: Expose CM4 eMMC as USB mass storage (skip if device already present)
if [ -n "${1:-}" ] && [ -b "${1}" ]; then
    echo "==> $1 already present — skipping rpiboot."
elif [ -z "${1:-}" ] && [ -n "$(find_removable_sd)" ]; then
    echo "==> Removable block device already present — skipping rpiboot."
else
    # NOTE: mass-storage-gadget64 is rejected by secure-boot-locked CM4s.
    #       Use secure-boot-msd (signed boot.img) instead.
    echo "==> Running rpiboot to expose CM4 eMMC (EMMC_DISABLE jumper must be bridged)..."
    sudo "$USBBOOT_DIR/rpiboot" -d "$MSD_DIR"
    echo "==> rpiboot done."
fi

# Step 2: Find the device (explicit arg or auto-detect)
if [ -n "${1:-}" ]; then
    DEVICE="$1"
    echo "==> Using device: $DEVICE — waiting for it to appear..."
    for i in $(seq 1 30); do
        [ -b "$DEVICE" ] && break
        sleep 1
        printf "  waiting for %s... (%ds)\r" "$DEVICE" "$i"
    done
    if [ ! -b "$DEVICE" ]; then
        echo "ERROR: $DEVICE did not appear as a block device within 30s."
        exit 1
    fi
else
    DEVICE=""
    echo "==> Waiting for removable block device..."
    for i in $(seq 1 30); do
        DEV=$(find_removable_sd)
        if [ -n "$DEV" ]; then
            DEVICE="/dev/$DEV"
            break
        fi
        sleep 1
        printf "  waiting... (%ds)\r" "$i"
    done
    if [ -z "$DEVICE" ]; then
        echo "ERROR: No removable block device found within 30s."
        echo "       Run 'lsblk' to find it, then re-run: $0 /dev/sdX"
        exit 1
    fi
    echo "==> Auto-detected CM4 eMMC at $DEVICE"
fi

# Step 3: Safety check - refuse to flash the host NVMe disk
if echo "$DEVICE" | grep -qE '^/dev/nvme'; then
    echo "ERROR: $DEVICE looks like the host NVMe. Aborting."
    exit 1
fi

# Step 4: Unmount any auto-mounted partitions
echo "==> Unmounting $DEVICE partitions..."
sudo umount "${DEVICE}"?* 2>/dev/null || true
sudo umount "${DEVICE}"[0-9]* 2>/dev/null || true

# Step 5: Flash via bmaptool
echo "==> Flashing $IMAGE -> $DEVICE ..."
sudo bmaptool copy "$IMAGE" "$DEVICE"
sudo sync

echo ""
echo "==> Flash complete!"
echo "    Remove the EMMC_DISABLE jumper, then power-cycle the CM4."