#!/bin/bash
# create-disk-image.sh — Create a raw disk image with boot + data partitions
# Phase 1: simple layout (boot + data). Phase 3 adds A/B system partitions.
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:-2048}"  # 2 GB default

VMLINUZ="$ROOTFS_DIR/vmlinuz"
INITRAMFS="$ROOTFS_DIR/kubesolo-os.gz"

for f in "$VMLINUZ" "$INITRAMFS"; do
    [ -f "$f" ] || { echo "ERROR: Missing $f — run 'make initramfs'"; exit 1; }
done

echo "==> Creating ${IMG_SIZE_MB}MB disk image..."
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: 256MB boot (ext4) + rest data (ext4)
# Using sfdisk for scriptability
sfdisk "$IMG_OUTPUT" << EOF
label: dos
unit: sectors

# Boot partition: 256 MB, bootable
start=2048, size=524288, type=83, bootable
# Data partition: remaining space
start=526336, type=83
EOF

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

cleanup() {
    umount "${LOOP}p1" 2>/dev/null || true
    umount "${LOOP}p2" 2>/dev/null || true
    losetup -d "$LOOP" 2>/dev/null || true
    rm -rf "$MNT_BOOT" "$MNT_DATA" 2>/dev/null || true
}
trap cleanup EXIT

# Format partitions
mkfs.ext4 -q -L KSOLOBOOT "${LOOP}p1"
mkfs.ext4 -q -L KSOLODATA "${LOOP}p2"

# Mount and populate boot partition
MNT_BOOT=$(mktemp -d)
MNT_DATA=$(mktemp -d)

mount "${LOOP}p1" "$MNT_BOOT"
mount "${LOOP}p2" "$MNT_DATA"

# Install syslinux + kernel + initramfs to boot partition
mkdir -p "$MNT_BOOT/boot/syslinux"
cp "$VMLINUZ" "$MNT_BOOT/boot/vmlinuz"
cp "$INITRAMFS" "$MNT_BOOT/boot/kubesolo-os.gz"

# Syslinux config for disk boot (extlinux)
cat > "$MNT_BOOT/boot/syslinux/syslinux.cfg" << 'EOF'
DEFAULT kubesolo
TIMEOUT 30
PROMPT 0

LABEL kubesolo
    KERNEL /boot/vmlinuz
    INITRD /boot/kubesolo-os.gz
    APPEND quiet kubesolo.data=LABEL=KSOLODATA

LABEL kubesolo-debug
    KERNEL /boot/vmlinuz
    INITRD /boot/kubesolo-os.gz
    APPEND kubesolo.data=LABEL=KSOLODATA kubesolo.debug console=ttyS0,115200n8

LABEL kubesolo-shell
    KERNEL /boot/vmlinuz
    INITRD /boot/kubesolo-os.gz
    APPEND kubesolo.shell console=ttyS0,115200n8
EOF

# Install extlinux bootloader
if command -v extlinux >/dev/null 2>&1; then
    extlinux --install "$MNT_BOOT/boot/syslinux" 2>/dev/null || {
        echo "WARN: extlinux install failed — image may not be directly bootable"
        echo "      Use with QEMU -kernel/-initrd flags instead"
    }
fi

# Prepare data partition structure
for dir in kubesolo containerd etc-kubesolo log usr-local network; do
    mkdir -p "$MNT_DATA/$dir"
done

sync

echo ""
echo "==> Disk image created: $IMG_OUTPUT"
echo "    Size: $(du -h "$IMG_OUTPUT" | cut -f1)"
echo "    Boot partition (KSOLOBOOT): kernel + initramfs"
echo "    Data partition (KSOLODATA): persistent K8s state"