#!/usr/bin/env python3
# SPDX-License-Identifier: Linux-OpenIB
# Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES
# PYTHON_ARGCOMPLETE_OK
from __future__ import annotations
import argparse
import collections
import importlib
import inspect
import itertools
import os
import re
import subprocess
import sys
import tempfile

from typing import *

BDF_RE = re.compile(r"^([0-9a-f]+?):([0-9a-f]{2}?):([0-9a-f]{2}?)\.([0-9a-f])$")
KERNEL_ACS_ISOLATED = "xx111x1"
pci_vendors = {
    "MELLANOX": 0x15B3,
    "NVIDIA": 0x10DE,
}


class CommandError(Exception):
    pass


def sysfs_read_str(sysdir: str, fn: str) -> str:
    """Read the entire content of a sysfs file to a string"""
    with open(os.path.join(sysdir, fn)) as F:
        return F.read().strip()


def sysfs_read_link(sysdir: str, fn: str) -> str:
    """Read a link in sysfs to an absolute path string"""
    return os.readlink(os.path.join(sysdir, fn))


def PCI_VDEVICE(vendor: str, device_id: int) -> re.Pattern:
    """Match a Vendor and device ID"""
    vendor_id = pci_vendors[vendor]
    return re.compile(rf"^pci:v{vendor_id:08X}d{device_id:08X}.*$")


def PCI_DEVICE_CLASS(cid: int) -> re.Pattern:
    """Match by exact programming class using the int coding from the kernel"""
    class_id = (cid >> 16) & 0xFF
    subclass_id = (cid >> 8) & 0xFF
    progif = cid & 0xFF
    return re.compile(rf"^pci:.*bc{class_id:02X}sc{subclass_id:02X}i{progif:02X}.*$")


def PCI_NVGPU() -> re.Pattern:
    """Match all NVIDIA GPUs"""
    vendor_id = pci_vendors["NVIDIA"]
    class_id = 0x03
    return re.compile(rf"^pci:v{vendor_id:08X}.*bc{class_id:02X}.*$")


# Table of modalias matches to the device_type string
pci_device_types = {
    PCI_VDEVICE("NVIDIA", 0x22B1): "grace_rp",  # NVIDIA Grace PCI Root Port Bridge
    PCI_VDEVICE("NVIDIA", 0x22B2): "grace_rp",  # NVIDIA Grace PCI Root Port Bridge
    PCI_VDEVICE("NVIDIA", 0x22B8): "grace_rp",  # NVIDIA Grace PCI Root Port Bridge
    PCI_VDEVICE("MELLANOX", 0x1021): "cx_nic",  # ConnectX-7
    PCI_VDEVICE("MELLANOX", 0x1023): "cx_nic",  # ConnectX-8
    PCI_VDEVICE("MELLANOX", 0xA2DC): "bf3_nic",  # BlueField-3
    PCI_VDEVICE("MELLANOX", 0x2100): "cx_dma",  # ConnectX-8 DMA Controller
    PCI_VDEVICE("MELLANOX", 0x197B): "bf3_switch",  # USP/DSP of a BF3 switch
    PCI_VDEVICE("MELLANOX", 0x197C): "cx_switch",  # USP/DSP of a CX switch
    PCI_DEVICE_CLASS(0x010802): "nvme",
    PCI_NVGPU(): "nvgpu",
}


class PCIBDF(
    collections.namedtuple("PCIBDF", ["segment", "bus", "device", "function"])
):
    """Bus Device Function for a PCI device"""

    def as_pci(self):
        return f"{self.segment}:{self.bus}:{self.device}.{self.function}"

    def __str__(self):
        return self.as_pci()

    def __repr__(self):
        return f"PCIBDF({self.segment}, {self.bus}, {self.device}, {self.function})"


def to_pcibdf(s: str) -> Optional[PCIBDF]:
    g = BDF_RE.match(s)
    if not g:
        return None
    return PCIBDF(*g.groups())


class PCIDevice(object):
    device_type = ""
    vpd_v3: str = None
    parent: PCIDevice = None
    lspci_data: str = None

    def __init__(self, sysdir: str, bdf: PCIBDF):
        self.sysdir = sysdir
        self.bdf = bdf
        try:
            self.iommu_group = int(
                os.path.split(sysfs_read_link(sysdir, "iommu_group"))[-1]
            )
        except FileNotFoundError:
            self.iommu_group = None

        try:
            self.numa_node = int(sysfs_read_str(sysdir, "numa_node"))
        except FileNotFoundError:
            self.numa_node = None

        self.modalias = sysfs_read_str(sysdir, "modalias")
        for k, v in pci_device_types.items():
            if k.match(self.modalias):
                self.device_type = v
                break

        sysdir = os.path.realpath(sysdir)
        parent = os.path.basename(os.path.dirname(sysdir))
        self.parent_bdf = to_pcibdf(parent)
        self.children: Set[PCIDevice] = set()

    def finish_loading(self):
        """Do more expensive parsing operations"""
        if self.device_type == "cx_nic" or self.device_type == "cx_dma":
            self.vpd_v3 = self.parse_vpd_v3()
        if "switch" in self.device_type or self.device_type == "grace_rp":
            self.has_acs = self.parse_has_acs()

    def iterdownstream(self) -> Generator[PCIDevice, None, None]:
        """Iterate over all downstream devices of this device recursively"""
        for pdev in self.children:
            yield pdev
            yield from pdev.iterdownstream()

    def iterfulltree(self):
        for pdev in self.iterupstream_path():
            if not pdev.parent:
                yield from pdev.iterdownstream()

    def iterupstream_path(self):
        """Iterate over each step along the upstream path from the devices
        parent to the root."""
        pdev = self.parent
        while pdev:
            yield pdev
            pdev = pdev.parent

    def __repr__(self):
        return f"PCIDevice({self.bdf})"

    def lspci(self):
        """Fetch the verbose output of lspci"""
        vpdfn = os.path.join(self.sysdir, "vpd")
        if os.path.exists(vpdfn) and not os.access(vpdfn, os.R_OK):
            raise CommandError(
                f"Need access to the PCI VPD information in {vpdfn}, are you root?"
            )

        if not self.lspci_data:
            self.lspci_data = subprocess.check_output(
                ["lspci", "-s", f"{self.bdf.as_pci()}", "-vv"]
            ).decode()
        return self.lspci_data

    def parse_vpd_v3(self):
        """Use lspci to parse the VPD and get the V3 UUID, this only works as
        root on non-secure boot systems."""
        g = re.search(
            r"Capabilities: \[.+?\] Vital Product Data$.*Read-only fields:$.*\[V3\] Vendor specific: (.*?)$.*End$",
            self.lspci(),
            re.DOTALL | re.MULTILINE,
        )
        if not g:
            return None
        return g.group(1)

    def parse_has_acs(self):
        """True if the device has an ACS capability"""
        return bool(
            re.search(
                r"Capabilities: \[.+?\] Access Control Services$",
                self.lspci(),
                re.DOTALL | re.MULTILINE,
            )
        )

    def parse_vpd_name(self):
        g = re.search(
            r"Capabilities: \[.+?\] Vital Product Data$.*Product Name: (.*?)$.*End$",
            self.lspci(),
            re.DOTALL | re.MULTILINE,
        )
        if not g:
            return None
        return g.group(1).strip()

    def read_config(self, regname: str):
        """Use setpci to read a register"""
        return int(
            subprocess.check_output(
                ["setpci", "-r", "-s", str(self.bdf), "ECAP_ACS+0x6.w"]
            )
            .decode()
            .strip(),
            16,
        )

    def get_subsystems(self):
        """Return a list of subsystem the PCI device is connected to"""
        res: Dict[str, Set[str]] = collections.defaultdict(set)
        for fn in os.listdir(self.sysdir):
            if fn in {"drm", "infiniband", "net", "nvme"}:
                res[fn].update(os.listdir(os.path.join(self.sysdir, fn)))
        return res


class NVCX_Complex(object):
    """Hold the related PCI functions together. A complex includes a CX PF, a CX
    DMA function, an GPU and related PCI switches in the DMA function
    segment."""

    def __init__(self, cx_pfs: Set[PCIDevice], cx_dma: PCIDevice, nvgpu: PCIDevice):
        self.cx_pfs = cx_pfs - {cx_dma}
        self.cx_pf = sorted(self.cx_pfs, key=lambda x: x.bdf)[0]
        self.cx_dma = cx_dma
        self.nvgpu = nvgpu

        # Identify the switch ports that are part of the shared path that
        # handles the P2P traffic
        self.shared_usp = self.__find_shared_usp()
        for pdev in self.cx_dma.iterupstream_path():
            if pdev in self.shared_usp.children:
                self.cx_dma_dsp = pdev
        for pdev in self.nvgpu.iterupstream_path():
            if pdev in self.shared_usp.children:
                self.nvgpu_dsp = pdev

        # There can be a NVMe device connected to the CX NIC as well. For NVMe
        # it is best to match with GPUs on the same socket, so a NUMA aware
        # approach would be fine, but also the GPU/NIC/NVMe could be
        # consistently paired based on the physical layout.
        self.nvmes: Set[PCIDevice] = set()
        for pdev in self.cx_pf.iterfulltree():
            if pdev.device_type == "nvme":
                self.nvmes.add(pdev)

    def __find_shared_usp(self) -> PCIDevice:
        """Find the USP that is shared by both devices, the immediate downstream
        bus is the point in the topology where P2P traffic will switch from an
        upstream to downstream direction."""
        common_path = set(self.cx_dma.iterupstream_path()).intersection(
            set(self.nvgpu.iterupstream_path())
        )
        assert common_path

        for pdev in self.cx_dma.iterupstream_path():
            if pdev in common_path:
                assert pdev.device_type == "cx_switch"
                for i in pdev.children:
                    assert i.device_type == "cx_switch"
                return pdev

    def get_subsystems(self):
        subsystems: Dict[str, Set[str]] = collections.defaultdict(set)
        for pdev in itertools.chain(self.cx_pfs, [self.nvgpu, self.cx_dma], self.nvmes):
            for k, v in pdev.get_subsystems().items():
                subsystems[k].update(v)
        return subsystems


def check_parent(pdev: PCIDevice, parent_type: str):
    if not pdev or not pdev.parent:
        return None
    if pdev.parent.device_type != parent_type:
        return None
    return pdev.parent


class PCITopo(object):
    """Load the PCI topology from sysfs and organize it"""

    def __init__(self):
        self.devices = self.__load_devices("/sys/bus/pci/devices/")
        self.has_cx_dma = any(
            pdev.device_type == "cx_dma" for pdev in self.devices.values()
        )
        if self.has_cx_dma:
            for pdev in self.devices.values():
                pdev.finish_loading()
            self.__build_topo()

    def __load_devices(self, sysdir: str):
        res: Dict[PCIBDF, PCIDevice] = {}
        for fn in os.listdir(sysdir):
            bdf = to_pcibdf(fn)
            if not bdf:
                continue
            assert bdf not in res
            res[bdf] = PCIDevice(os.path.join(sysdir, fn), bdf)
        return res

    def __get_nvcx_complex(self, cx_dma: PCIDevice):
        """Match the topology for the switch complex using a CX DMA function and a
        single GPU. It has two nested switches:

        RP --> SW -> CX_DMA
                  -> SW -> GPU
        """
        assert cx_dma.device_type == "cx_dma"
        if not cx_dma.vpd_v3:
            raise ValueError(f"CX DMA function {cx_dma} does not have a VPD V3 UUID")

        # The DMA and PF are matched using the UUID from the VPD
        cx_pfs = self.vpd_v3s.get(cx_dma.vpd_v3)
        if cx_pfs is None:
            raise ValueError(
                f"CX DMA function {cx_dma} does not have a matching PF, V3 UUID matching failed"
            )
            return None

        # Path from the DMA to the root port
        cx_dma_dsp = check_parent(cx_dma, "cx_switch")
        cx_usp = check_parent(cx_dma_dsp, "cx_switch")
        grace_rp = check_parent(cx_usp, "grace_rp")
        if not grace_rp:
            raise ValueError(
                f"CX DMA function {cx_dma} has an unrecognized upstream path"
            )

        # Path from the GPU to the root port
        nvgpus = [
            pdev for pdev in grace_rp.iterdownstream() if pdev.device_type == "nvgpu"
        ]
        if len(nvgpus) != 1:
            raise ValueError(f"CX DMA function {cx_dma} does not have a nearby GPU")
        nvgpu = nvgpus[0]
        nvgpu_dsp2 = check_parent(nvgpu, "cx_switch")
        nvgpu_usp2 = check_parent(nvgpu_dsp2, "cx_switch")
        nvgpu_dsp1 = check_parent(nvgpu_usp2, "cx_switch")
        if cx_usp != check_parent(nvgpu_dsp1, "cx_switch"):
            raise ValueError(
                f"CX DMA function {cx_dma} has an unrecognized upstream path from the GPU"
            )

        # Sanity check there is nothing unexpected in the topology
        alldevs = {
            cx_dma,
            cx_dma_dsp,
            cx_usp,
            nvgpu,
            nvgpu_dsp2,
            nvgpu_usp2,
            nvgpu_dsp1,
        }
        topodevs = set(grace_rp.iterdownstream())
        if alldevs != topodevs:
            raise ValueError(
                f"CX DMA function {cx_dma} has unexpected PCI devices in the topology"
            )
        return NVCX_Complex(cx_pfs, cx_dma, nvgpu)

    def __build_topo(self):
        """Collect cross-device information together and build the NVCX_Complex
        objects for the cx_dma functions"""
        self.vpd_v3s: Dict[str, Set[PCIDevice]] = collections.defaultdict(set)
        for pdev in self.devices.values():
            if pdev.parent_bdf:
                pdev.parent = self.devices[pdev.parent_bdf]
                pdev.parent.children.add(pdev)

            # Many PCI functions may share the same V3
            if pdev.vpd_v3:
                self.vpd_v3s[pdev.vpd_v3].add(pdev)

        self.nvcxs: List[NVCX_Complex] = []
        for pdev in self.devices.values():
            if pdev.device_type == "cx_dma":
                nvcx = self.__get_nvcx_complex(pdev)
                self.nvcxs.append(nvcx)
        self.nvcxs.sort(key=lambda x: x.cx_pf.bdf)

    def compute_acs(self):
        """Return a dictionary of PCI devices and the ACS mask the device should
        have"""
        acs: Dict[PCIDevice, str] = {}
        for nvcx in self.nvcxs:
            # For the DSP in the shared switch toward the CX8 DMA Direct interface:
            # Enable these bits:
            # bit-4 : ACS Upstream Forwarding
            # bit-3 : ACS P2P Completion Redirect
            # bit-0 : ACS Source Validation
            # Disable these bits:
            # bit-2 : ACS P2P Request Redirect
            assert nvcx.cx_dma_dsp.has_acs
            acs[nvcx.cx_dma_dsp] = "xx110x1"

            # For the DSP in the shared switch toward the GPU:
            # Enable the following bits:
            # bit-4 : ACS Upstream Forwarding
            # bit-2 : ACS P2P Request Redirect
            # bit-0 : ACS Source Validation
            # Disable the following bits:
            # bit-3 : ACS P2P Completion Redirect
            assert nvcx.nvgpu_dsp.has_acs
            acs[nvcx.nvgpu_dsp] = "xx101x1"

            # Disable ACS SV on the root port, this forces the entire segment
            # into one iommu_group and avoids kernel bugs building groups for
            # irregular ACS.
            for pdev in nvcx.cx_dma_dsp.iterupstream_path():
                if not pdev.parent:
                    assert pdev.has_acs
                    acs[pdev] = "xx111x0"

        # For all other CX bridges set kernel's default ACS enable
        # Enable these bits:
        # bit-4 : ACS Upstream Forwarding
        # bit-3 : ACS P2P Completion Redirect
        # bit-2 : ACS P2P Request Redirect
        # bit-0 : ACS Source Validation
        # Which match the kernel default
        for pdev in self.devices.values():
            if (
                pdev not in acs
                and ("switch" in pdev.device_type or "grace_rp" in pdev.device_type)
                and pdev.has_acs
            ):
                acs[pdev] = KERNEL_ACS_ISOLATED
        return acs


# -------------------------------------------------------------------
def print_list(title: str, items: list[str]):
    if not items:
        return
    if len(items) > 1:
        title = title + "s"
    list_str = ", ".join(sorted(items))
    print(f"\t{title}: {list_str}")


def args_topology(parser):
    parser.add_argument(
        "-j",
        "--json",
        action="store_true",
        dest="json",
        help="Output in machine readable JSON format",
    )


def topo_json(topo: PCITopo):
    import json

    jtop = []
    for nvcx in topo.nvcxs:
        jnvcx = {
            "rdma_nic_pf_bdf": str(nvcx.cx_pf.bdf),
            "rdma_dma_bdf": str(nvcx.cx_dma.bdf),
            "gpu_bdf": str(nvcx.nvgpu.bdf),
            "subsystems": {},
        }
        devname = nvcx.cx_pf.parse_vpd_name()
        if devname:
            jnvcx["rdma_nic_vpd_name"] = nvcx.cx_pf.parse_vpd_name()
        if nvcx.cx_pf.numa_node is not None:
            jnvcx["numa_node"] = nvcx.cx_pf.numa_node
        if nvcx.nvmes:
            jnvcx["nvme_bdf"] = str(next(iter(nvcx.nvmes)).bdf)

        for pdev in sorted(
            itertools.chain(nvcx.cx_pfs, [nvcx.nvgpu, nvcx.cx_dma], nvcx.nvmes),
            key=lambda x: x.bdf,
        ):
            subsys = pdev.get_subsystems()
            if subsys:
                jnvcx["subsystems"][str(pdev.bdf)] = {
                    subsys: list(devs) for subsys, devs in subsys.items()
                }
        jtop.append(jnvcx)
    print(json.dumps(jtop, indent=4))


def cmd_topology(args):
    """List the ConnectX NICs in the system with the corresponding NIC
    function, DMA Direct function and associated GPU."""
    topo = PCITopo()
    if not topo.has_cx_dma:
        raise CommandError("No ConnectX DMA Direct functions detected")

    if args.json:
        return topo_json(topo)

    for nvcx in topo.nvcxs:
        print(
            f"RDMA NIC={nvcx.cx_pf.bdf}, GPU={nvcx.nvgpu.bdf}, RDMA DMA Function={nvcx.cx_dma.bdf}"
        )

        devname = nvcx.cx_pf.parse_vpd_name()
        if devname:
            print(f"\t{devname}")

        if nvcx.cx_pf.numa_node is not None:
            print(f"\tNUMA Node: {nvcx.cx_pf.numa_node}")

        if len(nvcx.cx_pfs):
            print_list("NIC PCI device", [str(I.bdf) for I in nvcx.cx_pfs])

        subsystems = nvcx.get_subsystems()
        print_list("RDMA device", subsystems["infiniband"])
        print_list("Net device", subsystems["net"])
        print_list("DRM device", subsystems["drm"])
        print_list("NVMe device", subsystems["nvme"])
cmd_topology.__aliases__ = ("topo",)

# -------------------------------------------------------------------
def update_file(fn: str, new_content: str):
    """Make fn have new_content. If fn already has new_content nothing is
    done."""
    try:
        with open(fn, "rt") as F:
            old = F.read()
        if old == new_content:
            return False
    except FileNotFoundError:
        pass
    with tempfile.NamedTemporaryFile(dir=os.path.dirname(fn), mode="wt") as F:
        F.write(new_content)
        F.flush()
        os.chmod(F.name, 0o644)
        try:
            os.link(F.name, fn)
        except FileExistsError:
            os.unlink(fn)
            os.link(F.name, fn)
    return True


def args_write_grub_acs(parser):
    parser.add_argument(
        "-n",
        "--dry-run",
        action="store_true",
        dest="dry_run",
        help="Output the grub configuration to stdout and make no changes",
    )
    parser.add_argument(
        "--output",
        action="store",
        default="/etc/default/grub.d/config-acs.cfg",
        help="Grub dropin file to use for the kernel command line",
    )


def cmd_write_grub_acs(args):
    """Generate a grub dropin file to have the kernel commandline set the
    required ACS flags during system boot. This is the recommended way to
    configure ACS on systems but requires a compatible kernel.

    If the system does not have any need of ACS flags the dropin file will be
    removed. This command is intended for Debian style systems with a
    /etc/default/grub.d and update-grub command."""
    topo = PCITopo()
    if not topo.has_cx_dma:
        if args.dry_run:
            raise CommandError("No ConnectX DMA Direct functions detected")
        if os.path.exists(args.output):
            os.unlink(args.output)
        return

    acs = topo.compute_acs()
    config_acs = [
        f"{acs}@{pdev.bdf}"
        for pdev, acs in sorted(acs.items(), key=lambda x: x[0].bdf)
        if acs != KERNEL_ACS_ISOLATED
    ]
    acs_arg = ";".join(config_acs)
    grub_conf = [
        f"# Generated by {sys.argv[0]} do not change. ACS settings for RDMA GPU Direct",
        f'GRUB_CMDLINE_LINUX="$GRUB_CMDLINE_LINUX pci=config_acs=\\"{acs_arg}\\""',
    ]
    grub_conf = "\n".join(grub_conf)

    if args.dry_run:
        print(grub_conf)
        return

    try:
        os.makedirs(os.path.dirname(args.output))
    except FileExistsError:
        pass
    if update_file(args.output, grub_conf + "\n"):
        subprocess.check_call(["update-grub"])


# -------------------------------------------------------------------
def combine_acs(cur_acs, new_acs):
    for idx, val in enumerate(new_acs[::-1]):
        if val == "1":
            cur_acs = cur_acs | (1 << idx)
        elif val == "0":
            cur_acs = cur_acs & (0xFFFF ^ (1 << idx))
    return cur_acs


def args_setpci_acs(parser):
    parser.add_argument(
        "-n",
        "--dry-run",
        action="store_true",
        dest="dry_run",
        help="Output the setpci commands to stdout and make no changes",
    )


def cmd_setpci_acs(args):
    """Execute a series of set_pci commands that will immediately change the ACS
    settings to the required values. This is compatible with older kernels, but
    is not recommended. The kernel must boot with ACS enabled and the GPU driver
    must have the NVreg_GrdmaPciTopoCheckOverride=1 reg key set to disable
    safety checks that old kernels cannot support.

    NOTE: In this configuration unprivileged userspace can trigger platform RAS
    failures, use with caution!
    """
    topo = PCITopo()
    acs = topo.compute_acs()
    cmds: List[List[str]] = []
    for pdev, acs in sorted(acs.items(), key=lambda x: x[0].bdf):
        cur_acs = pdev.read_config("ECAP_ACS+0x6.w")
        new_acs = combine_acs(cur_acs, acs)
        if new_acs == cur_acs:
            continue

        cmd = ["setpci", "-r", "-s", str(pdev.bdf), f"ECAP_ACS+0x6.w={new_acs:04x}"]
        cmds.append(cmd)
    if args.dry_run:
        for cmd in cmds:
            print(" ".join(cmd))
        return
    for cmd in cmds:
        subprocess.check_call(cmd)


# -------------------------------------------------------------------
def args_check(parser):
    pass


def check_ok(msg: str):
    print(f"OK\t{msg}")


def check_fail(msg: str):
    print(f"FAIL\t{msg}")
    sys.exit(100)


def cmd_check(args):
    """Check that the running kernel and PCI environment are setup correctly for
    GPU Direct with ConnectX DMA Direct PCI functions."""
    topo = PCITopo()
    if not topo.has_cx_dma:
        raise CommandError("No ConnectX DMA Direct functions detected")
    check_ok("All ConnectX DMA functions have correct PCI topology")

    acs = topo.compute_acs()
    for pdev, acs in sorted(acs.items(), key=lambda x: x[0].bdf):
        cur_acs = pdev.read_config("ECAP_ACS+0x6.w")
        new_acs = combine_acs(cur_acs, acs)
        if new_acs == cur_acs:
            check_ok(
                f"ACS for {pdev.device_type} {pdev.bdf} has correct values {cur_acs:07b} = {acs}"
            )
        else:
            check_fail(
                f"ACS for {pdev.device_type} {pdev.bdf} has incorrect values {cur_acs:07b} != {acs}, (0x{cur_acs:x} != 0x{new_acs:x})"
            )

    # Correct iommu_groups are required to avoid NVreg_GrdmaPciTopoCheckOverride
    for nvcx in topo.nvcxs:
        if (
            nvcx.cx_dma.iommu_group == nvcx.nvgpu.iommu_group
            and nvcx.cx_dma.iommu_group is not None
        ):
            check_ok(
                f"Kernel iommu_group for DMA {nvcx.cx_dma.bdf} and GPU {nvcx.nvgpu.bdf} are both {nvcx.cx_dma.iommu_group}"
            )
        else:
            check_fail(
                f"Kernel iommu_group for DMA {nvcx.cx_dma.bdf} and GPU {nvcx.nvgpu.bdf} are not equal {nvcx.cx_dma.iommu_group} != {nvcx.nvgpu.iommu_group}"
            )


# -------------------------------------------------------------------
def load_all_commands(name):
    module = importlib.import_module(name)
    for k in dir(module):
        fn = getattr(module, k)
        argsfn = getattr(module, "args_" + k[4:], None)
        if argsfn is None or not k.startswith("cmd_") or not inspect.isfunction(fn):
            continue
        yield (k, fn, argsfn)


def get_cmd_aliases(fn):
    if hasattr(fn, "__aliases__"):
        return fn.__aliases__
    return ()

def main():
    parser = argparse.ArgumentParser(
        formatter_class=argparse.RawDescriptionHelpFormatter,
        description="""NVIDIA ConnectX GPU Direct ACS tool for Direct NIC platforms

This tool is used to view and control the PCI Access Control Flags (ACS) related
to the Direct NIC topology on supported NVIDIA platforms with ConnectX and
Blackwell family GPUs.

Direct NIC platforms have a unique multipath PCI topology where the ConnectX
has a main PCI function and a related DMA Direct function linked to the GPU.

This platform requires specific ACS flags in the PCI topology for reliable
operation, this tool helps users generate ACS settings for the local system.
""",
    )
    subparsers = parser.add_subparsers(title="Sub Commands", dest="command")
    subparsers.required = True

    commands = [I for I in load_all_commands(__name__)]
    commands.sort()

    # build sub parsers for all the loaded commands
    for k, fn, argsfn in commands:
        sparser = subparsers.add_parser(
            k[4:].replace("_", "-"), aliases=get_cmd_aliases(fn), help=fn.__doc__
        )
        sparser.required = True
        argsfn(sparser)
        sparser.set_defaults(func=fn)

    try:
        import argcomplete

        argcomplete.autocomplete(parser)
    except ImportError:
        pass

    # argparse will set 'func' to the cmd_* that executes this command
    args = parser.parse_args()
    try:
        args.func(args)
    except CommandError as e:
        print(f"E: {e}")
        sys.exit(100)


main()