#!/bin/bash
# Run PXE and/or ISO boot integration tests.
#
# Usage:
#   sudo ./scripts/test-provision.sh              # run PXE + ISO (x86_64)
#   sudo ./scripts/test-provision.sh pxe          # PXE only
#   sudo ./scripts/test-provision.sh iso          # ISO only (x86_64)
#   sudo ./scripts/test-provision.sh arm          # ARM ISO boot (emulated, SLOW ~60min)
#   sudo ./scripts/test-provision.sh arm-pxe      # ARM network PXE: NBP + discover + rescue (~20-30min)
#   sudo ./scripts/test-provision.sh arm-pxe-full # ARM network PXE incl. full install (~60-90min)
#   sudo ./scripts/test-provision.sh all          # all tests including ARM
#
# Prerequisites:
#   libvirtd, OVMF (edk2-ovmf), iPXE (ipxe-bootimgs-x86),
#   dnsmasq, xorriso, mtools, virt-install, qemu-img
#   ARM: qemu-system-aarch64, edk2-aarch64, ipxe-bootimgs-aarch64
set -e

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

cd "$PROJECT_ROOT"

# Detect real user for SSH keys
REAL_USER="${SUDO_USER:-$(whoami)}"
REAL_HOME=$(getent passwd "$REAL_USER" | cut -d: -f6)

RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BOLD='\033[1m'
RESET='\033[0m'

echo ""
echo -e "${BOLD}Lab Bastion -- Provision Integration Tests${RESET}"
echo "==========================================="
echo ""

# --- Prerequisite checks ---
MISSING=""
for cmd in virsh virt-install qemu-img dnsmasq xorriso mformat mcopy curl; do
    if ! command -v "$cmd" &>/dev/null; then
        MISSING="$MISSING $cmd"
    fi
done

if [ -n "$MISSING" ]; then
    echo -e "${RED}Missing tools:${RESET}$MISSING"
    echo "Install: sudo dnf install libvirt virt-install qemu-img dnsmasq xorriso mtools curl"
    exit 1
fi

if ! systemctl is-active libvirtd &>/dev/null; then
    echo -e "${RED}libvirtd not running.${RESET} Start with: sudo systemctl start libvirtd"
    exit 1
fi

if [ ! -f /usr/share/edk2/ovmf/OVMF_CODE.fd ]; then
    echo -e "${RED}OVMF firmware not found.${RESET} Install: sudo dnf install edk2-ovmf"
    exit 1
fi

MODE="${1:-both}"

# iPXE binaries are per-architecture. x86_64 is always required (the dnsmasq config
# references it); arm64 only for the ARM network-PXE modes.
IPXE_EFI=""
for f in /usr/share/ipxe/ipxe-snponly-x86_64.efi /usr/share/ipxe/ipxe-snp-x86_64.efi /usr/share/ipxe/ipxe-x86_64.efi; do
    [ -f "$f" ] && IPXE_EFI="$f" && break
done
if [ -z "$IPXE_EFI" ]; then
    echo -e "${RED}iPXE EFI binary not found.${RESET} Install: sudo dnf install ipxe-bootimgs-x86"
    exit 1
fi

IPXE_EFI_ARM64=""
for f in /usr/share/ipxe/arm64-efi/snponly.efi /usr/share/ipxe/arm64-efi/ipxe.efi; do
    [ -f "$f" ] && IPXE_EFI_ARM64="$f" && break
done

case "$MODE" in
    arm-pxe|arm-pxe-full|all)
        if [ -z "$IPXE_EFI_ARM64" ] && [ "$MODE" != "all" ]; then
            echo -e "${RED}arm64 iPXE binary not found.${RESET} Install: sudo dnf install ipxe-bootimgs-aarch64"
            exit 1
        fi
        ;;
esac

# Find SSH key
SSH_KEY=""
for name in id_ed25519 id_ecdsa id_rsa; do
    if [ -f "$REAL_HOME/.ssh/$name" ] && [ -f "$REAL_HOME/.ssh/$name.pub" ]; then
        SSH_KEY="$REAL_HOME/.ssh/$name"
        break
    fi
done
if [ -z "$SSH_KEY" ]; then
    echo -e "${RED}No SSH key found in $REAL_HOME/.ssh/${RESET}"
    exit 1
fi

echo -e "  User:    ${BOLD}$REAL_USER${RESET}"
echo -e "  SSH key: ${BOLD}$SSH_KEY${RESET}"
echo -e "  iPXE:    ${BOLD}$IPXE_EFI${RESET}"
echo -e "  iPXE a64:${BOLD} ${IPXE_EFI_ARM64:-not installed}${RESET}"
echo ""

require_arm_emulation() {
    if ! command -v qemu-system-aarch64 &>/dev/null; then
        echo -e "${RED}qemu-system-aarch64 not found.${RESET} Install: sudo dnf install qemu-system-aarch64 edk2-aarch64"
        exit 1
    fi
    if [ ! -f /usr/share/edk2/aarch64/QEMU_EFI.fd ]; then
        echo -e "${RED}AAVMF firmware not found.${RESET} Install: sudo dnf install edk2-aarch64"
        exit 1
    fi
}

run_test() {
    local name="$1" pattern="$2"
    echo ""
    echo -e "${YELLOW}━━━ Running $name test ━━━${RESET}"
    echo ""

    if SSH_KEY_PATH="$SSH_KEY" HOME="$REAL_HOME" \
       npx vitest run -c tests/integration/vitest.config.ts -t "$pattern" 2>&1; then
        echo ""
        echo -e "${GREEN}✔ $name test passed${RESET}"
        return 0
    else
        echo ""
        echo -e "${RED}✘ $name test failed${RESET}"
        return 1
    fi
}

FAILED=0

case "$MODE" in
    pxe)
        run_test "PXE boot" "PXE boot" || FAILED=1
        ;;
    iso)
        run_test "ISO boot" "ISO boot" || FAILED=1
        ;;
    arm|arm-iso)
        require_arm_emulation
        echo -e "${YELLOW}ARM emulation is ~10x slower than native. Expect 30-60 minutes.${RESET}"
        run_test "ARM ISO boot" "ARM ISO" || FAILED=1
        ;;
    arm-pxe)
        require_arm_emulation
        echo -e "${YELLOW}ARM emulation is ~10x slower than native. Expect 20-30 minutes.${RESET}"
        echo -e "${YELLOW}Covers option 93 -> arm64 NBP, discovery, and rescue over SSH.${RESET}"
        echo -e "${YELLOW}For the full install too, use: $0 arm-pxe-full${RESET}"
        run_test "ARM PXE rescue" "ARM PXE rescue" || FAILED=1
        ;;
    arm-pxe-full)
        require_arm_emulation
        echo -e "${YELLOW}ARM emulation is ~10x slower than native. Expect 60-90 minutes.${RESET}"
        ARM_PXE_FULL=1 run_test "ARM PXE (rescue + install)" "ARM PXE" || FAILED=1
        ;;
    both)
        run_test "PXE boot" "PXE boot" || FAILED=1
        run_test "ISO boot" "ISO boot" || FAILED=1
        ;;
    all)
        run_test "PXE boot" "PXE boot" || FAILED=1
        run_test "ISO boot" "ISO boot" || FAILED=1
        if command -v qemu-system-aarch64 &>/dev/null; then
            echo -e "${YELLOW}ARM emulation is ~10x slower than native.${RESET}"
            run_test "ARM ISO boot" "ARM ISO" || FAILED=1
            if [ -n "$IPXE_EFI_ARM64" ]; then
                run_test "ARM PXE rescue" "ARM PXE rescue" || FAILED=1
            else
                echo -e "${YELLOW}Skipping ARM PXE test (ipxe-bootimgs-aarch64 not installed)${RESET}"
            fi
        else
            echo -e "${YELLOW}Skipping ARM tests (qemu-system-aarch64 not installed)${RESET}"
        fi
        ;;
    *)
        echo "Usage: $0 [pxe|iso|arm|arm-pxe|arm-pxe-full|both|all]"
        exit 1
        ;;
esac

echo ""
if [ "$FAILED" -eq 0 ]; then
    echo -e "${GREEN}${BOLD}All provision tests passed.${RESET}"
else
    echo -e "${RED}${BOLD}Some tests failed.${RESET}"
    exit 1
fi