#!/usr/bin/python3

# Author: Christian Kastner <ckk@kvr.at>
# License: MIT


"""List all GPUs that debci would use in a test.

These are all GPUs referred to by --gpu in debci's backend configuration
variables, autopkgtest_args_<backend>.

If no --gpu specifications exist, then this will fall back to use all
VFIO-assigned GPUs (qemu-based backends) resp. regular GPUs (podman-
based backends).
"""


import argparse
import subprocess
import sys


def debci_config_get_string(option: str) -> str:
    """Gets a debci config option, interpreting its value as a string."""
    cmd = ["debci", "config", "-v", option]
    out = subprocess.check_output(cmd, text=True)
    return out.split("\n", maxsplit=1)[0]


def debci_config_get_list(option: str) -> list[str]:
    """Gets a debci config option, interpreting its value as a list."""
    return debci_config_get_string(option).split()


def find_system_gpus() -> list[tuple[str, bool]]:
    """Return list of (GPU slot, VFIO flag) pairs for this system."""
    gpus = []
    # 0300=VGA compatible controller, 0380=Display controller
    for vendevcla in ["::0300", "::0380"]:
        cmd = ["lspci", "-D", "-d", vendevcla]
        lines = subprocess.check_output(cmd, text=True).split("\n")
        for line in lines:
            if not line.strip():
                continue
            slot_id = line.split()[0]

            cmd = ["lspci", "-s", slot_id, "-vv"]
            dlines = subprocess.check_output(cmd, text=True).split("\n")
            for dline in dlines:
                prop, _, val = dline.strip().partition(":")
                if prop == "Kernel driver in use":
                    gpus.append((slot_id, val.strip() == "vfio-pci"))
                    break
    return gpus


def main() -> None:
    """Main"""
    parser = argparse.ArgumentParser(
        description="List all GPUs relevant to a debci configuration"
    )
    parser.add_argument(
        "--csv",
        action="store_true",
        help="Output GPUs comma-separated instead of one per line",
    )
    args = parser.parse_args()

    backend = debci_config_get_string("backend")
    # Backend names may contain a +, the arg variable names may not
    backend_args = f"autopkgtest_args_{backend.replace('+', '')}"
    match backend:
        case name if name.startswith("qemu"):
            autopkgtest_args = debci_config_get_list(backend_args)
        case name if name.startswith("podman"):
            autopkgtest_args = debci_config_get_list(backend_args)
        case _:
            sys.exit(10)

    config_parser = argparse.ArgumentParser()
    config_parser.add_argument("--gpu", action="append")
    config_args, _ = config_parser.parse_known_args(autopkgtest_args)
    if config_args.gpu:
        if args.csv:
            print(",".join(config_args.gpu))
        else:
            for gpu in config_args.gpu:
                print(gpu)
        return

    system_gpus = find_system_gpus()

    # config did not have --gpu options, so get them all
    match backend:
        case name if name.startswith("qemu"):
            gpus = [gpu for gpu, vfio in system_gpus if vfio]
        case name if name.startswith("podman"):
            gpus = [gpu for gpu, vfio in system_gpus if not vfio]
        case _:
            sys.exit(10)

    if args.csv:
        print(",".join(gpus))
    else:
        for gpu in gpus:
            print(gpu)


if __name__ == "__main__":
    main()
