#!/bin/bash
# ============================================================================
# create_disk.sh — Create a bootable UEFI disk image
#
# Usage: create_disk.sh <arch> <efi_binary> <kernel_elf> <output_img>
#
# Creates a GPT disk image with a FAT32 EFI System Partition containing:
#   - EFI/BOOT/BOOTx64.EFI (or BOOTAA64.EFI) — our UEFI loader
#   - kernel.elf — the kernel binary
# ============================================================================

set -e

ARCH="$1"
EFI_BIN="$2"
KERNEL_ELF="$3"
OUTPUT="$4"

if [ -z "$ARCH" ] || [ -z "$EFI_BIN" ] || [ -z "$KERNEL_ELF" ] || [ -z "$OUTPUT" ]; then
    echo "Usage: $0 <arch> <efi_binary> <kernel_elf> <output_img>"
    exit 1
fi

# Determine the default EFI boot path
if [ "$ARCH" = "x86_64" ]; then
    EFI_BOOT_NAME="BOOTX64.EFI"
elif [ "$ARCH" = "aarch64" ]; then
    EFI_BOOT_NAME="BOOTAA64.EFI"
else
    echo "Unknown architecture: $ARCH"
    exit 1
fi

echo "Creating disk image for $ARCH..."

# Image size: 64 MiB (more than enough)
IMG_SIZE=$((64 * 1024 * 1024))

# Create empty image
dd if=/dev/zero of="$OUTPUT" bs=1M count=64 status=none

# Create GPT partition table with one EFI System Partition
# Use sgdisk if available, otherwise fall back to parted
if command -v sgdisk &> /dev/null; then
    sgdisk --clear \
           --new=1:2048:131038 --typecode=1:ef00 --change-name=1:"EFI System" \
           "$OUTPUT" > /dev/null 2>&1
elif command -v parted &> /dev/null; then
    parted -s "$OUTPUT" mklabel gpt
    parted -s "$OUTPUT" mkpart "EFI System" fat32 1MiB 63MiB
    parted -s "$OUTPUT" set 1 esp on
else
    echo "ERROR: Need sgdisk or parted to create GPT partition"
    exit 1
fi

# Format the ESP partition as FAT32
# Extract partition offset (sector 2048 * 512 = 1 MiB)
PART_OFFSET=$((2048 * 512))
PART_SIZE=$(((131038 - 2048 + 1) * 512))

# Create a separate FAT32 filesystem image
FAT_IMG="${OUTPUT}.fat"
dd if=/dev/zero of="$FAT_IMG" bs=512 count=$((PART_SIZE / 512)) status=none
mkfs.vfat -F 32 "$FAT_IMG" > /dev/null 2>&1

# Copy files into the FAT filesystem using mtools
# Set up mtools config for this image
export MTOOLSRC="$(mktemp)"
echo "drive c: file=\"$FAT_IMG\" offset=0" > "$MTOOLSRC"

# Create EFI boot directory and copy files
mmd -i "$FAT_IMG" ::EFI
mmd -i "$FAT_IMG" ::EFI/BOOT
mcopy -i "$FAT_IMG" "$EFI_BIN" "::EFI/BOOT/$EFI_BOOT_NAME"
mcopy -i "$FAT_IMG" "$KERNEL_ELF" "::kernel.elf"

# List contents for verification
echo "ESP contents:"
mdir -i "$FAT_IMG" ::/ -/ 2>/dev/null || true

# Write the FAT image into the partition slot
dd if="$FAT_IMG" of="$OUTPUT" bs=512 seek=2048 conv=notrunc status=none

# Clean up
rm -f "$FAT_IMG" "$MTOOLSRC"

echo "Disk image created: $OUTPUT"