#!/bin/bash
# create-disk-image.sh — Create a raw disk image with A/B system partitions
#
# Partition layout (GPT):
#   Part 1: EFI/Boot    (256 MB, FAT32)  — GRUB + grubenv + A/B boot logic
#   Part 2: System A    (512 MB, ext4)   — vmlinuz + kubesolo-os.gz (active)
#   Part 3: System B    (512 MB, ext4)   — vmlinuz + kubesolo-os.gz (passive)
#   Part 4: Data        (remaining, ext4) — persistent K8s state
set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(cd "$SCRIPT_DIR/../.." && pwd)"
ROOTFS_DIR="${ROOTFS_DIR:-$PROJECT_ROOT/build/rootfs-work}"
OUTPUT_DIR="${OUTPUT_DIR:-$PROJECT_ROOT/output}"
VERSION="$(cat "$PROJECT_ROOT/VERSION")"
OS_NAME="kubesolo-os"

IMG_OUTPUT="$OUTPUT_DIR/${OS_NAME}-${VERSION}.img"
IMG_SIZE_MB="${IMG_SIZE_MB:-4096}"  # 4 GB default (larger for A/B)

VMLINUZ="$ROOTFS_DIR/vmlinuz"
INITRAMFS="$ROOTFS_DIR/kubesolo-os.gz"
GRUB_CFG="$PROJECT_ROOT/build/grub/grub.cfg"
GRUB_ENV_DEFAULTS="$PROJECT_ROOT/build/grub/grub-env-defaults"

for f in "$VMLINUZ" "$INITRAMFS" "$GRUB_CFG" "$GRUB_ENV_DEFAULTS"; do
    [ -f "$f" ] || { echo "ERROR: Missing $f"; exit 1; }
done

echo "==> Creating ${IMG_SIZE_MB}MB disk image with A/B partitions..."
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 (GPT):
#   Part 1: 256 MB EFI System Partition (FAT32)
#   Part 2: 512 MB System A (Linux filesystem)
#   Part 3: 512 MB System B (Linux filesystem)
#   Part 4: Remaining — Data (Linux filesystem)
sfdisk "$IMG_OUTPUT" << EOF
label: gpt

# EFI/Boot partition: 256 MB
start=2048, size=524288, type=C12A7328-F81F-11D2-BA4B-00A0C93EC93B, name="EFI"
# System A partition: 512 MB
size=1048576, type=0FC63DAF-8483-4772-8E79-3D69D8477DE4, name="SystemA"
# System B partition: 512 MB
size=1048576, type=0FC63DAF-8483-4772-8E79-3D69D8477DE4, name="SystemB"
# Data partition: remaining
type=0FC63DAF-8483-4772-8E79-3D69D8477DE4, name="Data"
EOF

# Set up loop device with partition mappings
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_EFI=$(mktemp -d)
MNT_SYSA=$(mktemp -d)
MNT_SYSB=$(mktemp -d)
MNT_DATA=$(mktemp -d)

cleanup() {
    umount "$MNT_EFI" 2>/dev/null || true
    umount "$MNT_SYSA" 2>/dev/null || true
    umount "$MNT_SYSB" 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_EFI" "$MNT_SYSA" "$MNT_SYSB" "$MNT_DATA" 2>/dev/null || true
}
trap cleanup EXIT

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

# Mount all partitions
mount "$P1" "$MNT_EFI"
mount "$P2" "$MNT_SYSA"
mount "$P3" "$MNT_SYSB"
mount "$P4" "$MNT_DATA"

# --- EFI/Boot Partition ---
echo "    Installing GRUB..."
mkdir -p "$MNT_EFI/EFI/BOOT"
mkdir -p "$MNT_EFI/boot/grub"

# Copy GRUB config
cp "$GRUB_CFG" "$MNT_EFI/boot/grub/grub.cfg"

# Create GRUB environment file from defaults
if command -v grub-editenv >/dev/null 2>&1; then
    GRUB_EDITENV=grub-editenv
elif command -v grub2-editenv >/dev/null 2>&1; then
    GRUB_EDITENV=grub2-editenv
else
    GRUB_EDITENV=""
fi

GRUBENV_FILE="$MNT_EFI/boot/grub/grubenv"

if [ -n "$GRUB_EDITENV" ]; then
    # Create grubenv with defaults
    "$GRUB_EDITENV" "$GRUBENV_FILE" create
    while IFS='=' read -r key value; do
        # Skip comments and empty lines
        case "$key" in
            '#'*|'') continue ;;
        esac
        "$GRUB_EDITENV" "$GRUBENV_FILE" set "$key=$value"
    done < "$GRUB_ENV_DEFAULTS"
    echo "    GRUB environment created with grub-editenv"
else
    # Fallback: write grubenv file manually (1024 bytes, padded with '#')
    echo "    WARN: grub-editenv not found — writing grubenv manually"
    {
        echo "# GRUB Environment Block"
        while IFS='=' read -r key value; do
            case "$key" in
                '#'*|'') continue ;;
            esac
            echo "$key=$value"
        done < "$GRUB_ENV_DEFAULTS"
    } > "$GRUBENV_FILE.tmp"
    # Pad to 1024 bytes (GRUB requirement)
    truncate -s 1024 "$GRUBENV_FILE.tmp"
    mv "$GRUBENV_FILE.tmp" "$GRUBENV_FILE"
fi

# Install GRUB EFI binary if available
if command -v grub-mkimage >/dev/null 2>&1; then
    grub-mkimage -O x86_64-efi -o "$MNT_EFI/EFI/BOOT/bootx64.efi" \
        -p /boot/grub \
        part_gpt ext2 fat normal linux echo all_video test search \
        search_fs_uuid search_label configfile loadenv \
        2>/dev/null || echo "    WARN: grub-mkimage failed — use QEMU -bios flag"
elif command -v grub2-mkimage >/dev/null 2>&1; then
    grub2-mkimage -O x86_64-efi -o "$MNT_EFI/EFI/BOOT/bootx64.efi" \
        -p /boot/grub \
        part_gpt ext2 fat normal linux echo all_video test search \
        search_fs_uuid search_label configfile loadenv \
        2>/dev/null || echo "    WARN: grub2-mkimage failed — use QEMU -bios flag"
else
    echo "    WARN: grub-mkimage not found — EFI boot image not created"
    echo "          Install grub2-tools or use QEMU -kernel/-initrd flags"
fi

# For BIOS boot: install GRUB i386-pc modules if available
if command -v grub-install >/dev/null 2>&1; then
    grub-install --target=i386-pc --boot-directory="$MNT_EFI/boot" \
        --no-floppy "$LOOP" 2>/dev/null || {
        echo "    WARN: BIOS GRUB install failed — EFI-only or use QEMU -kernel"
    }
elif command -v grub2-install >/dev/null 2>&1; then
    grub2-install --target=i386-pc --boot-directory="$MNT_EFI/boot" \
        --no-floppy "$LOOP" 2>/dev/null || {
        echo "    WARN: BIOS GRUB install failed — EFI-only or use QEMU -kernel"
    }
fi

# --- System A Partition (active) ---
echo "    Populating System A (active)..."
cp "$VMLINUZ" "$MNT_SYSA/vmlinuz"
cp "$INITRAMFS" "$MNT_SYSA/kubesolo-os.gz"
echo "$VERSION" > "$MNT_SYSA/version"

# --- System B Partition (passive, initially same as A) ---
echo "    Populating System B (passive)..."
cp "$VMLINUZ" "$MNT_SYSB/vmlinuz"
cp "$INITRAMFS" "$MNT_SYSB/kubesolo-os.gz"
echo "$VERSION" > "$MNT_SYSB/version"

# --- Data Partition ---
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 "==> Disk image created: $IMG_OUTPUT"
echo "    Size: $(du -h "$IMG_OUTPUT" | cut -f1)"
echo "    Part 1 (KSOLOEFI):  GRUB + A/B boot config"
echo "    Part 2 (KSOLOA):    System A — kernel + initramfs (active)"
echo "    Part 3 (KSOLOB):    System B — kernel + initramfs (passive)"
echo "    Part 4 (KSOLODATA): Persistent K8s state"
echo ""