#!/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 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
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

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

# 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 ""

# --- Determine which tests to run ---
MODE="${1:-both}"

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)
        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
        echo -e "${YELLOW}ARM emulation is ~10x slower than native. Expect 30-60 minutes.${RESET}"
        run_test "ARM ISO boot" "ARM ISO" || 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
        else
            echo -e "${YELLOW}Skipping ARM test (qemu-system-aarch64 not installed)${RESET}"
        fi
        ;;
    *)
        echo "Usage: $0 [pxe|iso|arm|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