#!/bin/bash
# create-rpi-image.sh — Create a raw disk image for Raspberry Pi SD card
#
# Partition layout (MBR):
#   Part 1: Boot Control  (32 MB, FAT32, label KSOLOCTL) — firmware + autoboot.txt
#   Part 2: Boot A        (256 MB, FAT32, label KSOLOA)  — kernel + DTBs + initramfs
#   Part 3: Boot B        (256 MB, FAT32, label KSOLOB)  — same as Boot A (initially identical)
#   Part 4: Data          (remaining of 2GB, ext4, label KSOLODATA)
#
# The RPi EEPROM loads start4.elf from partition 1 (KSOLOCTL).
# autoboot.txt on partition 1 redirects boot_partition to 2 (Boot A) or 3 (Boot B).
# The firmware then loads config.txt, kernel, and initramfs from the selected partition.
#
# MBR is required — GPT + autoboot.txt is not reliably supported on Pi 4.
#
# Usage: build/scripts/create-rpi-image.sh
set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(cd "$SCRIPT_DIR/../.." && pwd)"

# shellcheck source=../config/versions.env
. "$SCRIPT_DIR/../config/versions.env"

ROOTFS_DIR="${ROOTFS_DIR:-$PROJECT_ROOT/build/rootfs-work}"
OUTPUT_DIR="${OUTPUT_DIR:-$PROJECT_ROOT/output}"
CACHE_DIR="${CACHE_DIR:-$PROJECT_ROOT/build/cache}"
VERSION="$(cat "$PROJECT_ROOT/VERSION")"

IMG_OUTPUT="$OUTPUT_DIR/${OS_NAME}-${VERSION}.rpi.img"
IMG_SIZE_MB="${IMG_SIZE_MB:-2048}"  # 2 GB default

# ARM64 kernel (Image format, not bzImage)
KERNEL="${CACHE_DIR}/custom-kernel-arm64/Image"
INITRAMFS="${ROOTFS_DIR}/kubesolo-os.gz"
RPI_FIRMWARE_DIR="${CACHE_DIR}/rpi-firmware"

echo "==> Creating ${IMG_SIZE_MB}MB Raspberry Pi disk image..."

# --- Verify required files ---
MISSING=0
for f in "$KERNEL" "$INITRAMFS"; do
    if [ ! -f "$f" ]; then
        echo "ERROR: Missing $f"
        MISSING=1
    fi
done

if [ ! -d "$RPI_FIRMWARE_DIR" ]; then
    echo "ERROR: Missing RPi firmware directory: $RPI_FIRMWARE_DIR"
    echo "  Run 'make fetch' to download firmware blobs."
    MISSING=1
fi

if [ "$MISSING" = "1" ]; then
    echo ""
    echo "Required files:"
    echo "  Kernel:    $KERNEL (run 'make kernel-arm64')"
    echo "  Initramfs: $INITRAMFS (run 'make initramfs')"
    echo "  Firmware:  $RPI_FIRMWARE_DIR/ (run 'make fetch')"
    exit 1
fi

mkdir -p "$OUTPUT_DIR"

# --- Create sparse image ---
dd if=/dev/zero of="$IMG_OUTPUT" bs=1M count=0 seek="$IMG_SIZE_MB" 2>/dev/null

# --- Partition table (MBR) ---
# MBR is required for reliable RPi boot with autoboot.txt.
# GPT + autoboot.txt fails on many Pi 4 EEPROM versions.
# Part 1: Boot Control  32 MB FAT32 (firmware + autoboot.txt)
# Part 2: Boot A       256 MB FAT32 (kernel + initramfs + DTBs)
# Part 3: Boot B       256 MB FAT32 (kernel + initramfs + DTBs)
# Part 4: Data         remaining ext4
sfdisk "$IMG_OUTPUT" << EOF
label: dos

# Boot Control partition: 32 MB, FAT32 (type 0c = W95 FAT32 LBA)
start=2048, size=65536, type=c, bootable
# Boot A partition: 256 MB, FAT32
size=524288, type=c
# Boot B partition: 256 MB, FAT32
size=524288, type=c
# Data partition: remaining, Linux
type=83
EOF

# --- Set up loop device ---
LOOP=$(losetup --show -f "$IMG_OUTPUT")
echo "==> Loop device: $LOOP"

# Use kpartx for reliable partition device nodes (works in Docker/containers)
USE_KPARTX=false
if [ ! -b "${LOOP}p1" ]; then
    if command -v kpartx >/dev/null 2>&1; then
        kpartx -a "$LOOP"
        USE_KPARTX=true
        sleep 1
        LOOP_NAME=$(basename "$LOOP")
        P1="/dev/mapper/${LOOP_NAME}p1"
        P2="/dev/mapper/${LOOP_NAME}p2"
        P3="/dev/mapper/${LOOP_NAME}p3"
        P4="/dev/mapper/${LOOP_NAME}p4"
    else
        # Retry with -P flag
        losetup -d "$LOOP"
        LOOP=$(losetup --show -fP "$IMG_OUTPUT")
        sleep 1
        P1="${LOOP}p1"
        P2="${LOOP}p2"
        P3="${LOOP}p3"
        P4="${LOOP}p4"
    fi
else
    P1="${LOOP}p1"
    P2="${LOOP}p2"
    P3="${LOOP}p3"
    P4="${LOOP}p4"
fi

MNT_CTL=$(mktemp -d)
MNT_BOOTA=$(mktemp -d)
MNT_BOOTB=$(mktemp -d)
MNT_DATA=$(mktemp -d)

cleanup() {
    umount "$MNT_CTL" 2>/dev/null || true
    umount "$MNT_BOOTA" 2>/dev/null || true
    umount "$MNT_BOOTB" 2>/dev/null || true
    umount "$MNT_DATA" 2>/dev/null || true
    if [ "$USE_KPARTX" = true ]; then
        kpartx -d "$LOOP" 2>/dev/null || true
    fi
    losetup -d "$LOOP" 2>/dev/null || true
    rm -rf "$MNT_CTL" "$MNT_BOOTA" "$MNT_BOOTB" "$MNT_DATA" 2>/dev/null || true
}
trap cleanup EXIT

# --- Format partitions ---
mkfs.vfat -F 32 -n KSOLOCTL "$P1"
mkfs.vfat -F 32 -n KSOLOA "$P2"
mkfs.vfat -F 32 -n KSOLOB "$P3"
mkfs.ext4 -q -L KSOLODATA "$P4"

# --- Mount all partitions ---
mount "$P1" "$MNT_CTL"
mount "$P2" "$MNT_BOOTA"
mount "$P3" "$MNT_BOOTB"
mount "$P4" "$MNT_DATA"

# --- Boot Control Partition (KSOLOCTL) ---
# The RPi EEPROM loads start4.elf from partition 1.
# autoboot.txt tells the firmware which partition has config.txt + kernel.
echo "    Writing autoboot.txt + firmware to boot control partition..."
cat > "$MNT_CTL/autoboot.txt" << 'AUTOBOOT'
[all]
tryboot_a_b=1
boot_partition=2
[tryboot]
boot_partition=3
AUTOBOOT

# Copy firmware blobs — REQUIRED on partition 1 for EEPROM to boot
if ls "$RPI_FIRMWARE_DIR"/start*.elf 1>/dev/null 2>&1; then
    cp "$RPI_FIRMWARE_DIR"/start*.elf "$MNT_CTL/"
fi
if ls "$RPI_FIRMWARE_DIR"/fixup*.dat 1>/dev/null 2>&1; then
    cp "$RPI_FIRMWARE_DIR"/fixup*.dat "$MNT_CTL/"
fi
if [ -f "$RPI_FIRMWARE_DIR/bootcode.bin" ]; then
    cp "$RPI_FIRMWARE_DIR/bootcode.bin" "$MNT_CTL/"
fi

# --- Helper: populate a boot partition ---
populate_boot_partition() {
    local MNT="$1"
    local LABEL="$2"

    echo "    Populating $LABEL..."

    # config.txt — Raspberry Pi boot configuration
    cat > "$MNT/config.txt" << 'CFGTXT'
arm_64bit=1
kernel=kernel8.img
initramfs kubesolo-os.gz followkernel
enable_uart=1
gpu_mem=16
dtoverlay=disable-wifi
dtoverlay=disable-bt
CFGTXT

    # cmdline.txt — kernel command line
    # Note: must be a single line
    echo "console=serial0,115200 console=tty1 kubesolo.data=LABEL=KSOLODATA quiet" > "$MNT/cmdline.txt"

    # Copy kernel as kernel8.img (RPi 3/4/5 ARM64 convention)
    cp "$KERNEL" "$MNT/kernel8.img"

    # Copy initramfs
    cp "$INITRAMFS" "$MNT/kubesolo-os.gz"

    # Copy DTB overlays
    if [ -d "$RPI_FIRMWARE_DIR/overlays" ]; then
        cp -r "$RPI_FIRMWARE_DIR/overlays" "$MNT/"
    fi

    # Copy base DTBs (bcm2710-*, bcm2711-*, bcm2712-*)
    if ls "$RPI_FIRMWARE_DIR"/bcm27*.dtb 1>/dev/null 2>&1; then
        cp "$RPI_FIRMWARE_DIR"/bcm27*.dtb "$MNT/"
    fi

    # Write version marker
    echo "$VERSION" > "$MNT/version.txt"
}

# --- Boot A Partition (KSOLOA) ---
populate_boot_partition "$MNT_BOOTA" "Boot A (KSOLOA)"

# --- Boot B Partition (KSOLOB, initially identical) ---
populate_boot_partition "$MNT_BOOTB" "Boot B (KSOLOB)"

# --- Data Partition (KSOLODATA) ---
echo "    Preparing data partition..."
for dir in kubesolo containerd etc-kubesolo log usr-local network images; do
    mkdir -p "$MNT_DATA/$dir"
done

sync

echo ""
echo "==> Raspberry Pi disk image created: $IMG_OUTPUT"
echo "    Size: $(du -h "$IMG_OUTPUT" | cut -f1)"
echo "    Part 1 (KSOLOCTL):  Firmware + autoboot.txt (boot control)"
echo "    Part 2 (KSOLOA):    Boot A — kernel + initramfs + DTBs"
echo "    Part 3 (KSOLOB):    Boot B — kernel + initramfs + DTBs"
echo "    Part 4 (KSOLODATA): Persistent K8s state"
echo ""
echo "Write to SD card with:"
echo "    sudo dd if=$IMG_OUTPUT of=/dev/sdX bs=4M status=progress"
echo ""