#!/usr/bin/env python3

# Copyright: 2026 Hector CAO <hector.cao@canonical.com>
# SPDX-License-Identifier: GPL-3.0-or-later

import subprocess
import sys
from functools import wraps

# Source packages considered part of the virt stack.
_VIRT_SOURCES = {
    "qemu",
    "qemu-hwe",
    "edk2",
    "edk2-hwe",
    "seabios",
    "seabios-hwe",
    "libvirt",
    "libvirt-hwe",
}

# special cases where packages don't have a clear counterpart in the other variant
_SKIPPED = {
    "ubuntu-helper-virt-hwe",
    "qemu-efi-arm", # qemu-efi-arm is removed in edk2 (resolute) but still temporarily present in the archive
}

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 get_package_manual():
    """Retrieve lists of packages marked as manual.

    Returns:
        List of packages marked as manual.
    """
    manual_packages = set()

    # Get manually installed packages
    result = subprocess.run(
        ["apt-mark", "showmanual"],
        capture_output=True,
        text=True,
        check=False,
    )
    if result.returncode == 0:
        manual_packages = set(result.stdout.strip().splitlines())

    return manual_packages

def get_all_virt_installed_packages():
    """Get list of all installed packages with their source and version.

    Returns:
        variant: variant package (ubuntu-virt or ubuntu-virt-hwe)
        List of installed virt-related packages.
    """
    variant = None
    virt_packages = []

    mark_manual = get_package_manual()

    result = subprocess.run(
        ["dpkg-query", "-W", "-f=${binary:Package}\t${Version}\t${source:Package}\t${Status}\n"],
        capture_output=True,
        text=True,
        check=False,
    )
    for line in result.stdout.splitlines():
        parts = line.split("\t", 3)
        if len(parts) != 4:
            continue
        pkg, version, src, status = parts
        # only consider packages that are installed (status contains "installed")
        if "installed" not in status:
            continue
        # detect variant
        # dpkg-query might return package name with :arch suffix (e.g libvirt0:amd64)
        pkg_name = pkg.split(":")[0]
        if pkg_name == 'ubuntu-virt-hwe' or pkg_name == 'ubuntu-virt':
            if variant:
                print(f"[TEST] Warning: Multiple ubuntu-virt variants detected")
            else:
                variant = pkg_name
        # collect packages with virt-related sources
        if src in _VIRT_SOURCES:
            virt_packages.append((pkg, version, src, pkg_name in mark_manual))
    return variant, virt_packages


def counterpart_for(pkg_full: str, variant_before: str) -> str:
    pkg_name = pkg_full.split(":", 1)[0]
    pkg_arch = pkg_full.split(":", 1)[1] if ":" in pkg_full else None

    if variant_before == "ubuntu-virt":
        pkg_name = pkg_name + "-hwe"
    elif variant_before == "ubuntu-virt-hwe" and pkg_name.endswith("-hwe"):
        pkg_name = pkg_name[:-4]

    return pkg_name + (f":{pkg_arch}" if pkg_arch else "")


def package_exists_in_apt(pkg: str) -> bool:
    result = run(["apt-cache", "show", pkg], log_cmd=False, check=False, capture_output=True)
    return result.returncode == 0

def test_switch(from_variant: str) -> bool:

    variant_before, before = get_all_virt_installed_packages()

    # switch variant (show status before and after the switch).
    run(["ubuntu_virt_helper"])
    run(["ubuntu_virt_helper", "switch", "--", "-y"])
    run(["ubuntu_virt_helper"])

    variant_after, after = get_all_virt_installed_packages()

    print(f"[TEST] Package set before: {before}")
    print(f"[TEST] Package set after: {after}")

    assert variant_before == from_variant
    assert variant_after == "ubuntu-virt-hwe" if from_variant == "ubuntu-virt" else "ubuntu-virt"

    # check before == after
    for pkg_before in before:
        if pkg_before[0] in _SKIPPED:
            continue
        counterpart_name = counterpart_for(pkg_before[0], variant_before)
        counterpart = next((p for p in after if p[0] == counterpart_name), None)
        # [3] = auto/manual mark
        if counterpart is None or (counterpart[3] != pkg_before[3]) or (not package_exists_in_apt(counterpart[0])):
            assert False, f"[TEST] Package {pkg_before} is missing its counterpart {counterpart} after the switch"

    return True

@log_test
def test_helper_switch_qemu_utils():
    """
    Test that switching from ubuntu-virt to ubuntu-virt-hwe works as expected.
    Test the case where a simple package without complex dependencies is manually installed
    before the switch, and that the manual mark is preserved after the switch.
    """
    run(["apt-get", "install", "-y", "qemu-utils-hwe"])
    return test_switch("ubuntu-virt-hwe")

@log_test
def test_helper_switch_libvirt_daemon_driver_qemu():
    """
    Test that switching from ubuntu-virt to ubuntu-virt-hwe works as expected.
    Test the case where a complex package of the virt set with dependencies
    is manually installed before the switch, and that the manual mark is
    preserved after the switch.
    """
    run(["apt-get", "install", "-y", "libvirt-daemon-driver-qemu"])
    return test_switch("ubuntu-virt")

@log_test
def test_helper_switch_debvm():
    """
    Test that switching from ubuntu-virt to ubuntu-virt-hwe works as expected.
    Test the case where a third-party package of the virt set with dependencies
    is manually installed before the switch, and that the manual mark is
    preserved after the switch.
    """
    run(["apt-get", "install", "-y", "debvm"])
    # manual install of a package (to set to manual)
    run(["apt-get", "install", "-y", "qemu-system"])
    return test_switch("ubuntu-virt")

if __name__ == "__main__":
    run(["apt-get", "update"], capture_output=True)
    test_helper_switch_qemu_utils()
    test_helper_switch_libvirt_daemon_driver_qemu()
    test_helper_switch_debvm()
    sys.exit(0)