#!/usr/bin/env python3

"""Validate the machine types exposed by QEMU system emulators.

This test probes "qemu-system-<arch> -M help" for x86_64, s390x and ppc64
and verifies that the advertised machine types meet Ubuntu's expectations.
"""

import gzip
import os
import re
import shutil
import subprocess
import sys
from functools import wraps

# The Ubuntu release the machine types are created for
# Should be bumped every time we update the machine types for a new release
_CURRENT_RELEASE = "stonking"

# QEMU system emulators to probe for machine types.
_QEMU_X86_SYSTEM_BINARY = "qemu-system-x86_64"
_QEMU_S390X_SYSTEM_BINARY = "qemu-system-s390x"
_QEMU_PPC_SYSTEM_BINARY = "qemu-system-ppc64"

# Trailing annotations that "qemu-system-* -M help" appends to a description,
# e.g. "Standard PC (i440FX + PIIX, 1996) (alias of pc-i440fx-9.0)".
_ALIAS_RE = re.compile(r"\(alias of (?P<target>\S+)\)")
_DEFAULT_RE = re.compile(r"\(default\)")
_DEPRECATED_RE = re.compile(r"\(deprecated\)")

def log_test(test_func):
    @wraps(test_func)
    def wrapper(*args, **kwargs):
        print(f"[TEST] START: {test_func.__name__}")
        try:
            result = test_func(*args, **kwargs)
            print(f"[TEST] PASS: {test_func.__name__}")
            return result
        except AssertionError as exc:
            print(f"[TEST] FAIL: {test_func.__name__}: {exc}")
            raise
        except Exception as exc:
            print(f"[TEST] ERROR: {test_func.__name__}: {exc}")
            raise

    return wrapper


def run(cmd, log_cmd=True, check=True, capture_output=False):
    if log_cmd:
        print(f"[TEST] Running command: {' '.join(cmd)}")
    return subprocess.run(cmd, check=check, text=True, capture_output=capture_output)


def current_ubuntu_release():
    return _CURRENT_RELEASE

def parse_machine_types(output):
    """Parse the output of "qemu-system-* -M help" into a dictionary.

    The output looks like::

        Supported machines are:
        none                 empty machine
        pc                   Standard PC (...) (alias of pc-i440fx-9.0)
        pc-i440fx-9.0        Standard PC (...) (default)
        q35                  Standard PC (Q35 + ICH9, 2009)

    Args:
        output: The raw stdout produced by "qemu-system-* -M help".

    Returns:
        A dictionary keyed by machine name. Each value is a dict with:
            description: The human-readable description, annotations stripped.
            alias_of:    The target machine name if this is an alias, else None.
            is_default:  True if this machine is the default for the emulator.
            is_deprecated: True if this machine is flagged as deprecated.
    """
    machines = {}

    for line in output.splitlines():
        line = line.rstrip()
        # Skip the header and any blank lines.
        if not line or line.startswith("Supported machines are:"):
            continue
        # Name and description are separated by runs of whitespace.
        parts = line.split(None, 1)
        if not parts:
            continue
        name = parts[0]
        description = parts[1] if len(parts) > 1 else ""

        alias_match = _ALIAS_RE.search(description)
        alias_of = alias_match.group("target") if alias_match else None
        is_default = bool(_DEFAULT_RE.search(description))
        is_deprecated = bool(_DEPRECATED_RE.search(description))

        # Strip the trailing annotations to keep only the description text.
        description = _ALIAS_RE.sub("", description)
        description = _DEFAULT_RE.sub("", description)
        description = _DEPRECATED_RE.sub("", description)
        description = description.strip()

        machines[name] = {
            "description": description,
            "alias_of": alias_of,
            "is_default": is_default,
            "is_deprecated": is_deprecated,
        }

    return machines

def get_machine_types(binary):
    """Return the parsed machine-type dictionary for a qemu-system binary."""
    result = run(
        [binary, "-M", "help"],
        check=True,
        capture_output=True,
    )
    return parse_machine_types(result.stdout)


# Exactly one machine should be flagged as the default.
def check_one_default(machines):
    defaults = [name for name, info in machines.items() if info["is_default"]]
    assert len(defaults) == 1, f"expected 1 default, got {defaults}"

# Every alias must resolve to a machine that is listed.
def check_alias(machines):
    for name, info in machines.items():
        target = info["alias_of"]
        if target is not None:
            assert target in machines, (
                f"{_QEMU_SYSTEM_BINARY}: alias {name} points at unknown machine {target}"
            )

### S390x

# The default machine should be the ccw latest Ubuntu release
# s390-ccw-virtio-<release>[-v<integer>]
# and s390-ccw-virtio should be an alias of the default ccw machine.
def s390x_check_machines_ubuntu(machines):
    s390x_default_pattern = rf"s390-ccw-virtio-{current_ubuntu_release()}(-v\d+)?$"
    defaults = [name for name, info in machines.items() if info["is_default"]]
    assert defaults, "no default machine found"
    default_machine = defaults[0]
    assert re.match(s390x_default_pattern, default_machine), f"default machine {default_machine} does not match expected pattern"
    machines["s390-ccw-virtio"]["alias_of"] = default_machine

### PowerPC 64

# The default machine should be the pseries latest Ubuntu release
# pseries-<release>[-v<integer>]
# and pseries should be an alias of the default pseries machine.
def ppc_check_machines_ubuntu(machines):
    pseries_default_pattern = rf"pseries-{current_ubuntu_release()}(-v\d+)?$"
    defaults = [name for name, info in machines.items() if info["is_default"]]
    assert defaults, "no default machine found"
    default_machine = defaults[0]
    assert re.match(pseries_default_pattern, default_machine), f"default machine {default_machine} does not match expected pattern"
    # pseries should be an alias of the default pseries machine
    machines["pseries"]["alias_of"] = default_machine

### x86

# There should be an Ubuntu I440FX machine and it should be an alias to the current release I440FX Ubuntu machine.
# Example: ubuntu -> pc-i440fx-noble
def x86_check_ubuntu_i440fx(machines):
    ubuntu_i440fx_alias = "ubuntu"
    assert ubuntu_i440fx_alias in machines, f"{_QEMU_SYSTEM_BINARY}: machine {ubuntu_i440fx_alias} not found"
    target = machines[ubuntu_i440fx_alias]["alias_of"]
    # target should match the expected I440FX Ubuntu machine for x86.
    expected_target_pattern = rf"pc-i440fx-{current_ubuntu_release()}(-v\d+)?$"
    assert re.match(expected_target_pattern, target), f"{_QEMU_SYSTEM_BINARY}: alias {ubuntu_i440fx_alias} points at unexpected machine {target}"

# There should be an Ubuntu Q35 machine and it should be an alias to the current release Q35 Ubuntu machine.
# Example: ubuntu-q35 -> pc-q35-noble
def x86_check_ubuntu_q35(machines):
    ubuntu_q35_alias = "ubuntu-q35"
    assert ubuntu_q35_alias in machines, f"{_QEMU_SYSTEM_BINARY}: machine {ubuntu_q35_alias} not found"
    target = machines[ubuntu_q35_alias]["alias_of"]
    # target should match the expected Q35 Ubuntu machine for x86.
    expected_target_pattern = rf"pc-q35-{current_ubuntu_release()}(-v\d+)?$"
    assert re.match(expected_target_pattern, target), f"{_QEMU_SYSTEM_BINARY}: alias {ubuntu_q35_alias} points at unexpected machine {target}"

# The default machine should be the current release I440FX Ubuntu machine
# Expected name: pc-i440fx-<release>[-v<integer>]
def x86_check_ubuntu_default(machines):
    # Ensure that the default machine is an Ubuntu machine.
    defaults = [name for name, info in machines.items() if info["is_default"]]
    assert defaults, "no default machine found"
    default_machine = defaults[0]
    # check default_machine name pattern for x86 Ubuntu machines
    release = current_ubuntu_release()
    assert re.match(rf"pc-i440fx-{release}(-v\d+)?$", default_machine), f"default machine {default_machine} does not match expected pattern"

@log_test
def test_machine_types_s390x():
    machines = get_machine_types(_QEMU_S390X_SYSTEM_BINARY)
    assert machines, f"{_QEMU_S390X_SYSTEM_BINARY}: no machine types parsed"

    check_alias(machines)
    check_one_default(machines)

    s390x_check_machines_ubuntu(machines)

@log_test
def test_machine_types_ppc():
    machines = get_machine_types(_QEMU_PPC_SYSTEM_BINARY)
    assert machines, f"{_QEMU_PPC_SYSTEM_BINARY}: no machine types parsed"

    check_alias(machines)
    check_one_default(machines)

    ppc_check_machines_ubuntu(machines)

@log_test
def test_machine_types_x86():
    machines = get_machine_types(_QEMU_X86_SYSTEM_BINARY)
    assert machines, f"{_QEMU_X86_SYSTEM_BINARY}: no machine types parsed"

    check_alias(machines)
    check_one_default(machines)

    x86_check_ubuntu_i440fx(machines)
    x86_check_ubuntu_q35(machines)
    x86_check_ubuntu_default(machines)

    return True

if __name__ == "__main__":
    test_machine_types_x86()
    test_machine_types_s390x()
    test_machine_types_ppc()
    sys.exit(0)
