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

"""Common helpers shared by the hwe-* autopkgtest scripts.

This module is meant to be imported, not executed directly.
"""

import fnmatch
import os
import re
import subprocess
import sys

# Log file used by run_cmd/run_cmd_allow_fail; overridable by the caller by
# assigning hwe_lib.LOG_FILE before running commands.
LOG_FILE = os.environ.get("LOG_FILE", "hwe-check.log")

# Source packages that build the virt binary packages we manage.
# Base sources; each also has a matching "-hwe" source variant.
VIRT_BASE_SOURCES = ("qemu", "edk2", "seabios", "libvirt")
VIRT_BASE_SOURCE_REGEX = re.compile(r"^(" + "|".join(VIRT_BASE_SOURCES) + r")$")
VIRT_HWE_SOURCE_REGEX = re.compile(
    r"^(" + "|".join(f"{s}-hwe" for s in VIRT_BASE_SOURCES) + r")$"
)


class TestFailure(Exception):
    """Raised to abort a test with a message and a failure context."""

    def __init__(self, message, context):
        super().__init__(message)
        self.context = context


def step(message):
    print(f"[STEP] {message}")


def list_installed_pkgs_by_source(source_regex):
    """Return installed binary packages built from a matching source.

    Args:
        source_regex: A compiled regex matched against the source package name.

    Returns:
        A sorted list of installed binary package names.
    """
    result = subprocess.run(
        [
            "dpkg-query",
            "-W",
            "-f=${binary:Package}\t${db:Status-Status}\t${source:Package}\n",
        ],
        capture_output=True,
        text=True,
        check=False,
    )

    packages = set()
    for line in result.stdout.splitlines():
        parts = line.split("\t")
        if len(parts) != 3:
            continue
        pkg, status, src = parts
        if status != "installed":
            continue
        if not source_regex.match(src):
            continue
        if pkg == "ubuntu-helper-virt-hwe":
            continue
        packages.add(pkg)

    return sorted(packages)


def list_installed_virt_packages():
    """Return installed virt packages from base and -hwe sources, sorted."""
    packages = set(list_installed_pkgs_by_source(VIRT_BASE_SOURCE_REGEX))
    packages.update(list_installed_pkgs_by_source(VIRT_HWE_SOURCE_REGEX))
    return sorted(packages)


def run_cmd(description, *args):
    """Run a command, appending its output to LOG_FILE. Returns True on success."""
    step(description)
    with open(LOG_FILE, "a") as log:
        result = subprocess.run(
            list(args), stdout=log, stderr=subprocess.STDOUT, check=False
        )
    return result.returncode == 0


def run_cmd_allow_fail(description, *args):
    """Run a command, appending its output to LOG_FILE. Failures are ignored."""
    run_cmd(description, *args)


def systemd_check(*args):
    """Run 'systemctl <args>' quietly; return True on a zero exit status."""
    return (
        subprocess.run(
            ["systemctl", *args],
            stdout=subprocess.DEVNULL,
            stderr=subprocess.DEVNULL,
            check=False,
        ).returncode
        == 0
    )


def is_installed(package_name):
    result = subprocess.run(
        ["dpkg-query", "-W", "-f=${db:Status-Status}", package_name],
        capture_output=True,
        text=True,
        check=False,
    )
    return result.stdout.strip() == "installed"


def log_installed_virt_packages(context):
    packages = list_installed_virt_packages()

    step(f"  >>> Installed virt packages ({context}):")
    if not packages:
        step("  (none)")
        return

    for pkg in packages:
        step(f"  - {pkg}")


def exit_with_failure(message, context):
    print(f"ERROR: {message}")
    raise TestFailure(message, context)


def run_main(main_func):
    """Run a test entry point, printing a virt package snapshot on failure."""
    try:
        main_func()
    except TestFailure as exc:
        log_installed_virt_packages(f"failure snapshot ({exc.context})")
        sys.exit(1)
    except Exception:
        log_installed_virt_packages("failure snapshot (unexpected failure)")
        raise


def remove_all_virt_packages():
    packages = list_installed_virt_packages()

    if not packages:
        step("No virt packages installed; nothing to remove")
        return

    step("Running: apt-get remove -y " + " ".join(packages))
    run_cmd_allow_fail(
        "Removing all installed virt packages", "apt-get", "remove", "-y", *packages
    )


def remove_all_hwe_virt_packages():
    packages = list_installed_pkgs_by_source(VIRT_HWE_SOURCE_REGEX)

    if not packages:
        step("No -hwe virt packages installed; nothing to remove")
        return

    step("Running: apt-get remove -y " + " ".join(packages))
    run_cmd_allow_fail(
        "Removing all installed -hwe virt packages",
        "apt-get",
        "remove",
        "-y",
        *packages,
    )


def verify_only_variant(variant):
    """Verify that only the requested virt variant ('base' or 'hwe') is installed.

    Returns True if no package of the other variant is present and at least one
    package of the requested variant is installed.
    """
    base_pkgs = list_installed_pkgs_by_source(VIRT_BASE_SOURCE_REGEX)
    hwe_pkgs = list_installed_pkgs_by_source(VIRT_HWE_SOURCE_REGEX)

    if variant == "base":
        if hwe_pkgs:
            print(
                "ERROR [variant-check]: unexpected -hwe package(s) installed: "
                + " ".join(hwe_pkgs)
            )
            return False
        if not base_pkgs:
            print("ERROR [variant-check]: no base virt package installed")
            return False
    elif variant == "hwe":
        if base_pkgs:
            print(
                "ERROR [variant-check]: unexpected base package(s) installed: "
                + " ".join(base_pkgs)
            )
            return False
        if not hwe_pkgs:
            print("ERROR [variant-check]: no -hwe virt package installed")
            return False
    else:
        print(
            f"ERROR [variant-check]: unknown variant '{variant}' "
            "(expected 'base' or 'hwe')"
        )
        return False

    return True


def find_matching_entries(root, pattern):
    """Return filesystem entries under root whose basename matches pattern."""
    matches = []
    for dirpath, dirnames, filenames in os.walk(root):
        for name in dirnames + filenames:
            if fnmatch.fnmatch(name, pattern):
                matches.append(os.path.join(dirpath, name))
    return matches
