#!/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"
IMAGE="$REPO_ROOT/output/images/sdcard.img.xz"

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

# Step 1: Expose CM4 eMMC as USB mass storage
echo "==> Running rpiboot to expose CM4 eMMC (EMMC_DISABLE jumper must be bridged)..."
sudo "$USBBOOT_DIR/rpiboot" -d "$USBBOOT_DIR/mass-storage-gadget64"
echo "==> rpiboot done, waiting for block device..."

# Step 2: Find the device (explicit arg or auto-detect USB disk ~8 GiB)
if [ -n "${1:-}" ]; then
    DEVICE="$1"
    echo "==> Using specified device: $DEVICE"
else
    DEVICE=""
    for i in $(seq 1 30); do
        sleep 1
        # Detect USB block device of 7-8 GiB (CM4 eMMC)
        DEVICE=$(lsblk -dno NAME,TRAN,SIZE \
            | awk '$2=="usb" && ($3~/^7\.[0-9]+G$/ || $3~/^8\.[0-9]+G$/) {print "/dev/"$1}' \
            | head -1)
        [ -n "$DEVICE" ] && break
        printf "  waiting... (%ds)\r" "$i"
    done
    if [ -z "$DEVICE" ]; then
        echo "ERROR: CM4 eMMC did not appear as a USB block device within 30s."
        echo "       Run 'lsblk' manually and re-run with explicit device: $0 /dev/sdX"
        exit 1
    fi
    echo "==> Auto-detected CM4 eMMC at $DEVICE"
fi

# Step 3: Safety check - refuse to flash the host nvme/sata disk
if echo "$DEVICE" | grep -qE '^/dev/(nvme|sd[a-z]{2,}|sda$)'; then
    lsblk -dno TRAN "$DEVICE" | grep -qx usb || {
        echo "ERROR: $DEVICE does not appear to be a USB device. 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."