#!/bin/bash
# KNEL-Football ISO VM Testing Framework
# Uses libvirt/virsh to test ISO boot and runtime behavior
# Copyright © 2026 Known Element Enterprises LLC
# License: GNU Affero General Public License v3.0 only

set -euo pipefail

# Configuration
SCRIPT_DIR=""
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
readonly SCRIPT_DIR
readonly ISO_PATH="${SCRIPT_DIR}/output/knel-football-secure-v1.0.0.iso"
readonly VM_NAME="knel-football-test"
readonly VM_RAM="2048"
readonly VM_CPUS="2"
readonly VM_DISK_SIZE="10G"
readonly VM_DISK_PATH="/tmp/${VM_NAME}.qcow2"

# Colors for output
readonly RED='\033[0;31m'
readonly GREEN='\033[0;32m'
readonly YELLOW='\033[1;33m'
readonly NC='\033[0m' # No Color

# Logging functions
log_info() { echo -e "${GREEN}[INFO]${NC} $1"; }
log_warn() { echo -e "${YELLOW}[WARN]${NC} $1"; }
log_error() { echo -e "${RED}[ERROR]${NC} $1"; }

# Check prerequisites
check_prerequisites() {
    log_info "Checking prerequisites..."
    
    # Check for virsh command
    if ! command -v virsh &> /dev/null; then
        log_error "virsh command not found"
        log_error "Install libvirt: sudo apt install libvirt-clients libvirt-daemon-system qemu-system-x86"
        return 1
    fi
    
    # Check actual libvirt access (not just group membership)
    if ! virsh list &> /dev/null; then
        log_error "Cannot connect to libvirt"
        log_error "Ensure libvirtd is running and you have access"
        log_error "Try: sudo usermod -aG libvirt \$USER && logout/login"
        return 1
    fi
    
    # Check for qemu-img command
    if ! command -v qemu-img &> /dev/null; then
        log_error "qemu-img command not found"
        log_error "Install qemu: sudo apt install qemu-utils"
        return 1
    fi
    
    # Check ISO exists
    if [[ ! -f "$ISO_PATH" ]]; then
        log_error "ISO not found at: $ISO_PATH"
        log_error "Build the ISO first: ./run.sh iso"
        return 1
    fi
    
    log_info "All prerequisites satisfied"
    return 0
}

# Create VM disk image
create_disk() {
    log_info "Creating disk image: $VM_DISK_PATH"
    rm -f "$VM_DISK_PATH"
    qemu-img create -f qcow2 "$VM_DISK_PATH" "$VM_DISK_SIZE"
}

# Create and start VM
create_vm() {
    log_info "Creating VM: $VM_NAME"
    
    # Destroy existing VM if present
    virsh destroy "$VM_NAME" 2>/dev/null || true
    virsh undefine "$VM_NAME" 2>/dev/null || true
    
    # Create disk
    create_disk
    
    # Try virt-install first, fall back to direct QEMU
    if virt-install \
        --connect qemu:///session \
        --name "$VM_NAME" \
        --ram "$VM_RAM" \
        --vcpus "$VM_CPUS" \
        --disk path="$VM_DISK_PATH",format=qcow2 \
        --cdrom "$ISO_PATH" \
        --os-variant debian12 \
        --network user \
        --graphics vnc,listen=0.0.0.0 \
        --boot uefi \
        --noautoconsole \
        --virt-type kvm 2>&1; then
        log_info "VM created via virt-install"
    else
        log_warn "virt-install failed, using direct QEMU..."
        # Find UEFI firmware if available
        local uefi_fw=""
        for fw in /usr/share/OVMF/OVMF_CODE.fd /usr/share/qemu/OVMF_CODE.fd /usr/share/AAVMF/AAVMF_CODE.fd; do
            if [[ -f "$fw" ]]; then
                uefi_fw="$fw"
                break
            fi
        done
        
        local uefi_opts=()
        if [[ -n "$uefi_fw" ]]; then
            uefi_opts=(-bios "$uefi_fw")
            log_info "Using UEFI firmware: $uefi_fw"
        else
            log_warn "No UEFI firmware found, using legacy BIOS"
        fi
        
        # Use QEMU directly as fallback
        qemu-system-x86_64 \
            -name "$VM_NAME" \
            -m "$VM_RAM" \
            -smp "$VM_CPUS" \
            -drive file="$VM_DISK_PATH",format=qcow2,if=virtio \
            -cdrom "$ISO_PATH" \
            -netdev user,id=net0 \
            -device virtio-net-pci,netdev=net0 \
            -vnc :0 \
            "${uefi_opts[@]}" \
            -enable-kvm \
            -daemonize \
            -pidfile "/tmp/${VM_NAME}.pid"
        log_info "VM started via QEMU (PID: $(cat /tmp/${VM_NAME}.pid 2>/dev/null || echo 'unknown'))"
    fi
}

# Connect to VM console
connect_console() {
    log_info "Connecting to VM console..."
    virsh console "$VM_NAME"
}

# Get VM status
vm_status() {
    log_info "VM Status for: $VM_NAME"
    virsh dominfo "$VM_NAME" 2>/dev/null || log_error "VM not running"
}

# Check if VM is running
is_vm_running() {
    # Check virsh first
    if virsh domstate "$VM_NAME" 2>/dev/null | grep -q "running"; then
        return 0
    fi
    # Check QEMU pidfile
    if [[ -f "/tmp/${VM_NAME}.pid" ]]; then
        local pid
        pid=$(cat "/tmp/${VM_NAME}.pid")
        if kill -0 "$pid" 2>/dev/null; then
            return 0
        fi
    fi
    return 1
}

# Wait for boot and capture screenshot
capture_boot_screen() {
    local output_dir="${SCRIPT_DIR}/tmp/vm-screenshots"
    mkdir -p "$output_dir"
    
    log_info "Capturing boot screen..."
    virsh screenshot "$VM_NAME" "${output_dir}/boot-screen.ppm" 2>/dev/null || {
        log_warn "Could not capture screenshot"
    }
}

# Destroy VM and cleanup
destroy_vm() {
    log_info "Destroying VM: $VM_NAME"
    # Try virsh first
    virsh destroy "$VM_NAME" 2>/dev/null || true
    virsh undefine "$VM_NAME" 2>/dev/null || true
    # Kill QEMU process if running
    if [[ -f "/tmp/${VM_NAME}.pid" ]]; then
        local pid
        pid=$(cat "/tmp/${VM_NAME}.pid")
        kill "$pid" 2>/dev/null || true
        rm -f "/tmp/${VM_NAME}.pid"
    fi
    rm -f "$VM_DISK_PATH"
    log_info "Cleanup complete"
}

# Run automated boot test
run_boot_test() {
    log_info "Running automated boot test..."
    
    if ! check_prerequisites; then
        return 1
    fi
    
    create_vm
    
    log_info "Waiting for VM to boot (30 seconds)..."
    sleep 30
    
    if is_vm_running; then
        log_info "VM is running - boot test PASSED"
        vm_status
        capture_boot_screen
        return 0
    else
        log_error "VM not running - boot test FAILED"
        return 1
    fi
}

# Test Secure Boot
test_secure_boot() {
    log_info "Testing Secure Boot..."
    
    if ! is_vm_running; then
        log_error "VM not running, start it first"
        return 1
    fi
    
    # Check if VM booted with Secure Boot
    # This is a basic check - more sophisticated checks would require
    # parsing the boot logs or using expect scripts
    log_info "Secure Boot verification requires manual console inspection"
    log_info "Use: $0 console"
    log_info "Then check: dmesg | grep -i secure"
}

# Test FDE passphrase prompt
test_fde_prompt() {
    log_info "Testing FDE passphrase prompt..."
    
    if ! is_vm_running; then
        log_error "VM not running, start it first"
        return 1
    fi
    
    # FDE prompt appears during boot, requires console access
    log_info "FDE prompt verification requires manual console inspection"
    log_info "Use: $0 console"
    log_info "Watch for 'Please unlock disk' prompt during boot"
}

# Show usage
usage() {
    cat <<EOF
KNEL-Football ISO VM Testing Framework

Usage: $0 <command>

Commands:
  check       Check prerequisites (libvirt, ISO, etc.)
  create      Create and start test VM
  console     Connect to VM console
  status      Show VM status
  destroy     Destroy VM and cleanup
  boot-test   Run automated boot test
  secure-boot Test Secure Boot (manual verification)
  fde-test    Test FDE passphrase prompt (manual verification)
  help        Show this help message

Prerequisites:
  - User must be in libvirt group
  - libvirtd service must be running
  - ISO must exist in output/

Examples:
  $0 check          # Verify environment is ready
  $0 boot-test      # Create VM, boot ISO, verify boot
  $0 console        # Connect to running VM
  $0 destroy        # Clean up test VM

Note: After adding user to libvirt group, logout and login again.
EOF
    exit 1
}

# Main entry point
main() {
    local command="${1:-help}"
    
    case "$command" in
        check)
            check_prerequisites
            ;;
        create)
            check_prerequisites && create_vm
            ;;
        console)
            connect_console
            ;;
        status)
            vm_status
            ;;
        destroy)
            destroy_vm
            ;;
        boot-test)
            run_boot_test
            ;;
        secure-boot)
            test_secure_boot
            ;;
        fde-test)
            test_fde_prompt
            ;;
        help|*)
            usage
            ;;
    esac
}

main "$@"