#!/bin/bash
# dev-vm-arm64.sh — Launch ARM64 QEMU VM for development
#
# Two modes:
#
#   Default (direct kernel boot — fast iteration):
#     qemu loads the kernel Image + initramfs directly via -kernel/-initrd.
#     Skips bootloader, UEFI firmware, and disk image entirely.
#     Use this for kernel and init-script changes.
#
#   --disk (full UEFI boot — integration testing):
#     qemu boots the .arm64.img disk image via UEFI firmware -> GRUB -> kernel.
#     Exercises the full boot chain. Use this when changing the disk image
#     layout, GRUB config, or anything that touches the EFI partition.
#
# Usage:
#   ./hack/dev-vm-arm64.sh                       # direct kernel boot (default)
#   ./hack/dev-vm-arm64.sh --disk                # full UEFI boot from built image
#   ./hack/dev-vm-arm64.sh --debug               # enable kubesolo.debug
#   ./hack/dev-vm-arm64.sh --shell               # drop to emergency shell
#   ./hack/dev-vm-arm64.sh --disk /path/to.img   # boot a specific disk image
#   ./hack/dev-vm-arm64.sh <kernel> <initramfs>  # direct boot with custom files
set -euo pipefail

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

MODE="kernel"           # kernel | disk
VMLINUZ=""
INITRD=""
DISK_IMAGE=""
EXTRA_APPEND=""

while [ $# -gt 0 ]; do
    case "$1" in
        --shell)   EXTRA_APPEND="$EXTRA_APPEND kubesolo.shell"; shift ;;
        --debug)   EXTRA_APPEND="$EXTRA_APPEND kubesolo.debug"; shift ;;
        --disk)
            MODE="disk"
            shift
            # Optional next-arg as disk image path
            if [ $# -gt 0 ] && [ -f "$1" ]; then
                DISK_IMAGE="$1"
                shift
            fi
            ;;
        *)
            if [ "$MODE" = "kernel" ] && [ -z "$VMLINUZ" ]; then
                VMLINUZ="$1"
            elif [ "$MODE" = "kernel" ] && [ -z "$INITRD" ]; then
                INITRD="$1"
            fi
            shift
            ;;
    esac
done

# ---------------------------------------------------------------------------
# UEFI firmware probe (used for --disk mode)
# ---------------------------------------------------------------------------
find_uefi_firmware() {
    local candidates=(
        /usr/share/qemu-efi-aarch64/QEMU_EFI.fd
        /usr/share/AAVMF/AAVMF_CODE.fd
        /usr/share/edk2/aarch64/QEMU_EFI.fd
        /usr/share/qemu/edk2-aarch64-code.fd
        /opt/homebrew/share/qemu/edk2-aarch64-code.fd
        /usr/local/share/qemu/edk2-aarch64-code.fd
    )
    for f in "${candidates[@]}"; do
        [ -f "$f" ] && echo "$f" && return 0
    done
    return 1
}

# ---------------------------------------------------------------------------
# mkfs.ext4 probe (kernel mode creates a scratch data disk)
# ---------------------------------------------------------------------------
find_mkfs_ext4() {
    if command -v mkfs.ext4 >/dev/null 2>&1; then
        echo "mkfs.ext4"
    elif [ -x "/opt/homebrew/opt/e2fsprogs/sbin/mkfs.ext4" ]; then
        echo "/opt/homebrew/opt/e2fsprogs/sbin/mkfs.ext4"
    elif [ -x "/usr/local/opt/e2fsprogs/sbin/mkfs.ext4" ]; then
        echo "/usr/local/opt/e2fsprogs/sbin/mkfs.ext4"
    fi
}

# ===========================================================================
# Disk mode: boot the built .arm64.img through UEFI firmware + GRUB
# ===========================================================================
if [ "$MODE" = "disk" ]; then
    DISK_IMAGE="${DISK_IMAGE:-$PROJECT_ROOT/output/kubesolo-os-${VERSION}.arm64.img}"

    if [ ! -f "$DISK_IMAGE" ]; then
        echo "ERROR: Disk image not found: $DISK_IMAGE"
        echo "  Run 'make disk-image-arm64' to build it."
        exit 1
    fi

    UEFI_FW="$(find_uefi_firmware || true)"
    if [ -z "$UEFI_FW" ]; then
        echo "ERROR: No ARM64 UEFI firmware found."
        echo "  Install one of:"
        echo "    apt install qemu-efi-aarch64        # Debian/Ubuntu"
        echo "    dnf install edk2-aarch64            # Fedora/RHEL"
        echo "    brew install qemu                   # macOS (bundled)"
        exit 1
    fi

    # Pad UEFI firmware variable store to 64 MiB if QEMU expects pflash sizing.
    # Most ARM64 EFI .fd files are 64 MB; if yours is smaller, QEMU may refuse.
    echo "==> Launching ARM64 QEMU (UEFI disk boot)..."
    echo "    Firmware: $UEFI_FW"
    echo "    Disk:     $DISK_IMAGE"
    echo ""
    echo "    K8s API: localhost:6443"
    echo "    SSH:     localhost:2222"
    echo "    Press Ctrl+A X to exit QEMU"
    echo ""

    # -cpu max enables all emulated ARMv8 features (atomics, crypto, fp16).
    # piCore64's BusyBox is built with -march=armv8-a+crypto+lse and segfaults
    # under -cpu cortex-a72 because some required extensions aren't on by
    # default in that model.
    qemu-system-aarch64 \
        -machine virt \
        -cpu max \
        -m 2048 \
        -smp 2 \
        -nographic \
        -bios "$UEFI_FW" \
        -drive "file=$DISK_IMAGE,format=raw,if=virtio,media=disk" \
        -net "nic,model=virtio" \
        -net "user,hostfwd=tcp::6443-:6443,hostfwd=tcp::2222-:22"
    exit 0
fi

# ===========================================================================
# Kernel mode (default): direct -kernel / -initrd, fast iteration
# ===========================================================================
VMLINUZ="${VMLINUZ:-$PROJECT_ROOT/build/cache/kernel-arm64-generic/Image}"
INITRD="${INITRD:-$PROJECT_ROOT/build/rootfs-work/kubesolo-os.gz}"

# Fallback: previous-generation RPi kernel cache, in case someone hasn't yet
# rebuilt under v0.3 paths.
if [ ! -f "$VMLINUZ" ] && [ -f "$PROJECT_ROOT/build/cache/custom-kernel-rpi/Image" ]; then
    VMLINUZ="$PROJECT_ROOT/build/cache/custom-kernel-rpi/Image"
    echo "==> Note: falling back to RPi kernel ($VMLINUZ)"
fi

if [ ! -f "$VMLINUZ" ]; then
    echo "ERROR: Kernel not found: $VMLINUZ"
    echo "  Run 'make kernel-arm64' (generic) or 'make kernel-rpi' to build a kernel."
    exit 1
fi
if [ ! -f "$INITRD" ]; then
    echo "ERROR: Initrd not found: $INITRD"
    echo "  Run 'make rootfs-arm64' to build the initramfs."
    exit 1
fi

MKFS_EXT4="$(find_mkfs_ext4)"
if [ -z "$MKFS_EXT4" ]; then
    echo "ERROR: mkfs.ext4 not found. Install e2fsprogs:"
    if [ "$(uname)" = "Darwin" ]; then
        echo "  brew install e2fsprogs"
    else
        echo "  apt install e2fsprogs   # Debian/Ubuntu"
        echo "  dnf install e2fsprogs   # Fedora/RHEL"
    fi
    exit 1
fi

DATA_DISK="$(mktemp /tmp/kubesolo-arm64-data-XXXXXX).img"
dd if=/dev/zero of="$DATA_DISK" bs=1M count=1024 2>/dev/null
"$MKFS_EXT4" -q -L KSOLODATA "$DATA_DISK" 2>/dev/null
trap 'rm -f "$DATA_DISK"' EXIT

echo "==> Launching ARM64 QEMU (direct kernel boot)..."
echo "    Kernel: $VMLINUZ"
echo "    Initrd: $INITRD"
echo "    Data:   $DATA_DISK"
echo ""
echo "    K8s API: localhost:6443"
echo "    SSH:     localhost:2222"
echo "    Press Ctrl+A X to exit QEMU"
echo ""

qemu-system-aarch64 \
    -machine virt \
    -cpu max \
    -m 2048 \
    -smp 2 \
    -nographic \
    -kernel "$VMLINUZ" \
    -initrd "$INITRD" \
    -append "console=ttyAMA0 kubesolo.data=/dev/vda kubesolo.debug $EXTRA_APPEND" \
    -drive "file=$DATA_DISK,format=raw,if=virtio" \
    -net "nic,model=virtio" \
    -net "user,hostfwd=tcp::6443-:6443,hostfwd=tcp::2222-:22"