#!/usr/bin/env python3
"""Reproducible QEMU integration tests for danos.

For each test case it builds the kernel with `-Dtest-case=<name>`, boots it
headless in QEMU with the serial port captured to a file, and asserts that the
expected marker appears in that output before a timeout. The kernel's serial log
is machine-readable (see src/kernel/tests.zig and src/kernel/arch/*/serial.zig), so no
screen-scraping is involved.

Structured per-architecture so a second CPU (e.g. an AArch64 Raspberry Pi) is a
matter of adding an ARCHES entry and its serial/boot support — the test cases and
the runner stay the same.

Usage:
    python3 test/qemu_test.py                # run all cases on the default arch
    python3 test/qemu_test.py --arch x86_64  # pick an architecture
    python3 test/qemu_test.py smoke fault-pf # run only named cases
"""

import argparse
import os
import re
import shutil
import subprocess
import sys
import time

REPO = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
WORK = os.path.join(REPO, "zig-out", "qemu-test")

# --- Architecture configurations -------------------------------------------
# Each describes how to build and boot that architecture. Only the fields that
# actually differ between architectures live here.
ARCHES = {
    "x86_64": {
        "zig_flags": [],  # build.zig currently pins the kernel to x86_64
        "qemu": "qemu-system-x86_64",
        # Firmware lives in different places per OS/distro; resolve_firmware()
        # picks the first that exists so the harness runs unconfigured on Arch,
        # Debian/Ubuntu, Fedora, and macOS (Homebrew). Override by reordering or
        # dropping an absolute path at the front.
        "ovmf_code": [
            "/usr/share/edk2/x64/OVMF_CODE.4m.fd",          # Arch
            "/usr/share/OVMF/OVMF_CODE_4M.fd",              # Debian/Ubuntu
            "/usr/share/OVMF/OVMF_CODE.fd",                 # older Debian/Ubuntu
            "/usr/share/edk2-ovmf/x64/OVMF_CODE.fd",        # Fedora
            "/opt/homebrew/share/qemu/edk2-x86_64-code.fd", # macOS Homebrew (Apple Silicon)
            "/usr/local/share/qemu/edk2-x86_64-code.fd",    # macOS Homebrew (Intel)
        ],
        "ovmf_vars": [
            "/usr/share/edk2/x64/OVMF_VARS.4m.fd",          # Arch
            "/usr/share/OVMF/OVMF_VARS_4M.fd",              # Debian/Ubuntu
            "/usr/share/OVMF/OVMF_VARS.fd",                 # older Debian/Ubuntu
            "/usr/share/edk2-ovmf/x64/OVMF_VARS.fd",        # Fedora
            "/opt/homebrew/share/qemu/edk2-i386-vars.fd",   # macOS Homebrew (Apple Silicon)
            "/usr/local/share/qemu/edk2-i386-vars.fd",      # macOS Homebrew (Intel)
        ],
        "efi_app": ("EFI/BOOT/BOOTX64.efi", "BOOTX64.efi"),  # (dest in ESP, name in zig-out/bin)
        "kernel": ("kernel", "kernel"),
        # Further files shipped on the ESP: the init user program.
        "extra": [("sbin/init", "init")],
        # Built as a function so we can splice in per-run paths.
        "qemu_args": lambda a, esp, vars_fd, serial: [
            "-machine", "q35", "-m", "128M",
            "-drive", f"if=pflash,format=raw,readonly=on,file={a['ovmf_code']}",
            "-drive", f"if=pflash,format=raw,file={vars_fd}",
            "-drive", f"format=raw,file=fat:rw:{esp}",
            "-net", "none",
            "-vga", "none", "-device", "VGA,edid=on,xres=1280,yres=720",
            "-display", "none",
            "-serial", f"file:{serial}",
            "-no-reboot",  # exit instead of rebooting on a triple fault
        ],
    },
    # To add an architecture, e.g. "aarch64": provide its qemu binary, firmware,
    # boot method, and the AArch64 kernel/serial support in src/kernel/arch/aarch64/.
}

# --- Test cases ------------------------------------------------------------
# `expect`: a regex that must appear in serial output => pass.
# `fail`:   optional regex whose appearance => immediate fail.
CASES = [
    {"name": "smoke",
     "expect": r"DANOS-TEST-RESULT: PASS",
     "fail":   r"DANOS-TEST-RESULT: FAIL"},
    {"name": "discovery",
     "expect": r"DANOS-TEST-RESULT: PASS",
     "fail":   r"DANOS-TEST-RESULT: FAIL"},
    {"name": "wx",
     "expect": r"DANOS-TEST-RESULT: PASS",
     "fail":   r"DANOS-TEST-RESULT: FAIL"},
    {"name": "timer",
     "expect": r"DANOS-TEST-RESULT: PASS",
     "fail":   r"DANOS-TEST-RESULT: FAIL"},
    {"name": "clock",
     "expect": r"DANOS-TEST-RESULT: PASS",
     "fail":   r"DANOS-TEST-RESULT: FAIL"},
    {"name": "vmm",
     "expect": r"DANOS-TEST-RESULT: PASS",
     "fail":   r"DANOS-TEST-RESULT: FAIL"},
    {"name": "heap",
     "expect": r"DANOS-TEST-RESULT: PASS",
     "fail":   r"DANOS-TEST-RESULT: FAIL"},
    {"name": "sched",
     "expect": r"DANOS-TEST-RESULT: PASS",
     "fail":   r"DANOS-TEST-RESULT: FAIL"},
    {"name": "priority",
     "expect": r"DANOS-TEST-RESULT: PASS",
     "fail":   r"DANOS-TEST-RESULT: FAIL"},
    {"name": "sleep",
     "expect": r"DANOS-TEST-RESULT: PASS",
     "fail":   r"DANOS-TEST-RESULT: FAIL"},
    {"name": "event",
     "expect": r"DANOS-TEST-RESULT: PASS",
     "fail":   r"DANOS-TEST-RESULT: FAIL"},
    {"name": "ipc",
     "expect": r"DANOS-TEST-RESULT: PASS",
     "fail":   r"DANOS-TEST-RESULT: FAIL"},
    # Parallelism: needs more than one core, so this case boots with -smp 4.
    {"name": "smp",
     "smp":    4,
     "expect": r"DANOS-TEST-RESULT: PASS",
     "fail":   r"DANOS-TEST-RESULT: FAIL"},
    # Affinity: a pinned task must never migrate off its core.
    {"name": "affinity",
     "smp":    4,
     "expect": r"DANOS-TEST-RESULT: PASS",
     "fail":   r"DANOS-TEST-RESULT: FAIL"},
    # Stress the big kernel lock across cores; heavier, so a longer timeout.
    {"name": "smp-stress",
     "smp":    4,
     "timeout": 90,
     "expect": r"DANOS-TEST-RESULT: PASS",
     "fail":   r"DANOS-TEST-RESULT: FAIL"},
    # Retry: a forced first-wake failure must still bring every core online.
    {"name": "smp-retry",
     "smp":    4,
     "expect": r"DANOS-TEST-RESULT: PASS",
     "fail":   r"DANOS-TEST-RESULT: FAIL"},
    {"name": "fault-ud", "expect": r"invalid opcode \(vector 6\)"},
    {"name": "fault-pf", "expect": r"page fault \(vector 14\)"},
    {"name": "fault-df", "expect": r"double fault \(vector 8\)"},
    # Fault on an application processor: a #DF pinned to core 1 must be caught by that
    # core's own IST (N >= 1), proving per-core TSS works; a broken one triple-faults.
    {"name": "fault-ap-df",
     "smp":    4,
     "expect": r"core [1-9]\d*: double fault \(vector 8\)",
     "fail":   r"could not pin"},
    {"name": "fault-nx",
     "expect": r"page fault \(vector 14\)",
     "fail":   r"NX not enforced"},
    {"name": "fault-null", "expect": r"page fault \(vector 14\)"},
    # Ring 3: a user program runs at CPL 3, makes int 0x80 syscalls, survives
    # timer interrupts, and exits back into the kernel.
    {"name": "user",
     "expect": r"DANOS-TEST-RESULT: PASS",
     "fail":   r"DANOS-TEST-RESULT: FAIL"},
    # Isolation: a ring-3 read of a kernel-only page must #PF with error code
    # 0x5 (present|user) at the user RIP. ([\s\S] spans lines; `.` doesn't.)
    {"name": "user-pf",
     "expect": r"page fault \(vector 14\)[\s\S]*error code : 0x5[\s\S]*RIP\s*: 0x00007000000000",
     "fail":   r"DANOS-TEST-RESULT: FAIL"},
    # The real user binary: the bootloader ships sbin/init off the ESP, the
    # kernel loads the ELF and runs it in ring 3, and it writes + exits cleanly.
    {"name": "init",
     "expect": r"DANOS-TEST-RESULT: PASS",
     "fail":   r"DANOS-TEST-RESULT: FAIL"},
    # The ACPI power path succeeds by QEMU *exiting* (S5 off / reset), so match the
    # pre-transition marker; the FAIL line only appears if the transition didn't take.
    {"name": "poweroff",
     "expect": r"DANOS-POWER: attempting poweroff",
     "fail":   r"DANOS-TEST-RESULT: FAIL"},
    {"name": "reboot",
     "expect": r"DANOS-POWER: attempting reboot",
     "fail":   r"DANOS-TEST-RESULT: FAIL"},
]

TIMEOUT = 30  # seconds per case


def build(arch, case):
    cmd = ["zig", "build", f"-Dtest-case={case}"] + arch["zig_flags"]
    r = subprocess.run(cmd, cwd=REPO, capture_output=True, text=True)
    if r.returncode != 0:
        return r.stderr.strip() or r.stdout.strip()
    return None


def make_esp(arch):
    """Assemble a fresh EFI System Partition from the freshly built binaries."""
    esp = os.path.join(WORK, "esp")
    if os.path.exists(esp):
        shutil.rmtree(esp)
    efi_dest, efi_name = arch["efi_app"]
    kern_dest, kern_name = arch["kernel"]
    os.makedirs(os.path.join(esp, os.path.dirname(efi_dest)), exist_ok=True)
    shutil.copy(os.path.join(REPO, "zig-out", "bin", efi_name), os.path.join(esp, efi_dest))
    shutil.copy(os.path.join(REPO, "zig-out", "bin", kern_name), os.path.join(esp, kern_dest))
    for dest, name in arch.get("extra", []):
        os.makedirs(os.path.join(esp, os.path.dirname(dest)), exist_ok=True)
        shutil.copy(os.path.join(REPO, "zig-out", "bin", name), os.path.join(esp, dest))
    return esp


def resolve_firmware(arch):
    """Collapse the ovmf_code/ovmf_vars candidate lists to the first path that
    exists on this machine. Mutates `arch` in place; idempotent (a resolved
    string is left untouched). Exits with a clear message if none is found."""
    for key, label in (("ovmf_code", "OVMF_CODE"), ("ovmf_vars", "OVMF_VARS")):
        val = arch[key]
        if isinstance(val, str):
            continue  # already resolved on a previous call
        for cand in val:
            if os.path.exists(cand):
                arch[key] = cand
                break
        else:
            sys.exit(f"error: no {label} firmware image found; looked in:\n  "
                     + "\n  ".join(val)
                     + "\nInstall OVMF (edk2-ovmf / ovmf) or add its path above.")


def run_case(arch, case):
    err = build(arch, case["name"])
    if err:
        return False, "build failed:\n" + err

    esp = make_esp(arch)
    vars_fd = os.path.join(WORK, "vars.fd")
    shutil.copy(arch["ovmf_vars"], vars_fd)
    serial = os.path.join(WORK, "serial.log")
    if os.path.exists(serial):
        os.remove(serial)

    expect = re.compile(case["expect"])
    fail = re.compile(case["fail"]) if case.get("fail") else None

    cmd = [arch["qemu"]] + arch["qemu_args"](arch, esp, vars_fd, serial)
    if case.get("smp"):  # some cases need more than one core (e.g. parallelism)
        cmd += ["-smp", str(case["smp"])]
    qemu = subprocess.Popen(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    try:
        timeout = case.get("timeout", TIMEOUT)
        deadline = time.monotonic() + timeout
        while time.monotonic() < deadline:
            time.sleep(0.2)
            text = ""
            if os.path.exists(serial):
                with open(serial, "r", errors="replace") as f:
                    text = f.read()
            if fail and fail.search(text):
                return False, "hit failure marker"
            if expect.search(text):
                return True, "matched " + repr(case["expect"])
            if qemu.poll() is not None:  # QEMU exited on its own
                if expect.search(text):
                    return True, "matched " + repr(case["expect"])
                return False, "QEMU exited before matching (triple fault?)"
        return False, f"timed out after {timeout}s without matching {case['expect']!r}"
    finally:
        if qemu.poll() is None:
            qemu.terminate()
            try:
                qemu.wait(timeout=5)
            except subprocess.TimeoutExpired:
                qemu.kill()


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--arch", default="x86_64", choices=sorted(ARCHES))
    ap.add_argument("cases", nargs="*", help="case names to run (default: all)")
    args = ap.parse_args()

    arch = ARCHES[args.arch]
    if not shutil.which(arch["qemu"]):
        print(f"error: {arch['qemu']} not found on PATH", file=sys.stderr)
        return 2
    resolve_firmware(arch)

    selected = [c for c in CASES if not args.cases or c["name"] in args.cases]
    os.makedirs(WORK, exist_ok=True)

    print(f"danos qemu tests  arch={args.arch}  cases={len(selected)}\n")
    failures = 0
    for case in selected:
        print(f"  {case['name']:<12} ... ", end="", flush=True)
        ok, detail = run_case(arch, case)
        print(("PASS" if ok else "FAIL") + f"  ({detail})")
        if not ok:
            failures += 1

    print()
    total = len(selected)
    print(f"{total - failures}/{total} passed")
    return 1 if failures else 0


if __name__ == "__main__":
    sys.exit(main())