#!/usr/bin/env bash
#
# scripts/build-image.sh — assemble a bootable rs-dos disk image (UEFI/GPT)
#
# Layout:
#   Partition 1: ESP (FAT32) — rsboot bootloader + kernel
#
# Output:
#   rs-dos.img (or custom path) — bootable raw disk image for UEFI
#
# Usage:
#   ./scripts/build-image.sh [output.img]
#
set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
PROJECT_DIR="$(dirname "$SCRIPT_DIR")"

cd "$PROJECT_DIR"

BUILD_DIR="$PROJECT_DIR/build"
ESP_DIR="$PROJECT_DIR/esp"

IMAGE="${1:-$PROJECT_DIR/rs-dos.img}"
TARGET=x86_64-unknown-uefi
PROFILE=debug
OUT_DIR="target/${TARGET}/${PROFILE}"

BOOTLOADER="$BUILD_DIR/bootloader/BOOTX64.EFI"
KERNEL="$OUT_DIR/rsdos.efi"
BZIMAGE="$ESP_DIR/bzImage"
INITRD="$ESP_DIR/initramfs.cpio.gz"
BOOT_EFI="$ESP_DIR/EFI/BOOT/BOOTX64.EFI"

ESP_START=2048 # 1 MiB offset (standard alignment)
ESP_SIZE_MB=64

IMAGE_SIZE_MB=$((ESP_SIZE_MB + 2)) # ESP + GPT gap + backup

# ---- checks ----
for cmd in sfdisk mkfs.vfat mcopy truncate dd; do
	if ! command -v "$cmd" &>/dev/null; then
		echo "Error: $cmd not found"
		exit 1
	fi
done

# ---- build kernel ----
echo "[1/4] Building kernel..."
cargo build --package rsdos --target x86_64-unknown-uefi ${CARGO_FLAGS:-}

# ---- build bootloader ----
echo "[2/4] Building bootloader..."
if [ ! -f "$BOOTLOADER" ]; then
	"$SCRIPT_DIR/build-bootloader.sh"
fi

# ---- assemble ESP ----
echo "[3/4] Assembling ESP..."
mkdir -p "$ESP_DIR/EFI/BOOT"
cp "$KERNEL" "$BZIMAGE"
cp "$BOOTLOADER" "$BOOT_EFI"
if [ ! -f "$INITRD" ]; then
	printf '' | gzip >"$INITRD"
fi

# ---- create disk image ----
echo "[4/4] Writing disk image..."
rm -f "$IMAGE"
truncate -s "${IMAGE_SIZE_MB}M" "$IMAGE"

# GPT partition table: single EFI System partition
sfdisk --quiet "$IMAGE" <<EOF
label: gpt
type=U, start=$ESP_START, size=+
EOF

# Build FAT32 image with ESP contents
ESP_IMG=$(mktemp)
truncate -s "${ESP_SIZE_MB}M" "$ESP_IMG"
mkfs.vfat -F 32 -n "RS-DOS" "$ESP_IMG" >/dev/null 2>&1

# Copy files into FAT
mmd -i "$ESP_IMG" ::/EFI 2>/dev/null || true
mmd -i "$ESP_IMG" ::/EFI/BOOT 2>/dev/null || true
mcopy -i "$ESP_IMG" "$BOOT_EFI" ::/EFI/BOOT/BOOTX64.EFI
mcopy -i "$ESP_IMG" "$BZIMAGE" ::/bzImage 2>/dev/null || true
mcopy -i "$ESP_IMG" "$INITRD" ::/initramfs.cpio.gz 2>/dev/null || true

# Inject FAT into image at partition offset
dd if="$ESP_IMG" of="$IMAGE" bs=512 seek="$ESP_START" conv=notrunc 2>/dev/null
rm -f "$ESP_IMG"

echo ""
echo "=== Done: $IMAGE ==="
ls -lh "$IMAGE"
echo ""
echo "Boot with:  ./scripts/boot-image.sh"
echo "USB:        sudo dd if=$IMAGE of=/dev/sdX bs=4M status=progress"