#!/usr/bin/env python3
# Copyright (c) Advanced Micro Devices, Inc. All rights reserved.
#
# SPDX-License-Identifier: MIT

# pylint: disable=invalid-name

"""
Check if the necessary AMD Infinity Storage (AIS) components are installed
and whether at least one mounted volume can use hipFile's fast path.

If all required components are installed and at least one volume is
hipFile-capable the program exits with exit code 0. Otherwise it exits with a
non-zero exit code.

A volume is reported as hipFile-capable when it mirrors the fast-path
eligibility the hipFile library enforces for a regular file:

  * filesystem is xfs, or ext4 mounted with data=ordered (the ext4 default), and
  * direct I/O (O_DIRECT) opens succeed on the filesystem.

See src/amd_detail/mountinfo.cpp (fstype/journaling classification) and
src/amd_detail/backend/fastpath.cpp (the accept rule) for the authoritative
logic the volume scan reproduces.

The BACKING column (nvme, lvm, md, ...) is informational only: a local NVMe
(non multipath) device is required by amdgpu/kfd, but that constraint is
enforced by the driver, not by the library's fast-path scoring.
"""

# pylint: enable=invalid-name

import argparse
import ctypes
import ctypes.util
import errno
import glob
import gzip
import os
import subprocess
import sys

# Mounts backed by these (or any non-block source) can never use the fast path.
PSEUDO_FSTYPES = {
    "autofs",
    "bpf",
    "cgroup",
    "cgroup2",
    "configfs",
    "debugfs",
    "devpts",
    "devtmpfs",
    "efivarfs",
    "fuse.gvfsd-fuse",
    "fusectl",
    "hugetlbfs",
    "mqueue",
    "overlay",
    "proc",
    "pstore",
    "ramfs",
    "securityfs",
    "squashfs",
    "sysfs",
    "tmpfs",
    "tracefs",
}


# amdgpu DKMS builds newer than this expose an authoritative AIS P2P-DMA state
# through the KFD topology "capability" bit. On such drivers we trust that bit
# instead of the kernel's CONFIG_PCI_P2PDMA, which only reports build-time
# support, not whether AIS actually came up.
KFD_AIS_CAPABILITY_MIN_VERSION = 2327936

# Bit set in a KFD topology node's "capability" mask once AIS P2P-DMA is
# initialized for that GPU.
KFD_CAPABILITY_AIS = 0x40

# Per-GPU topology directory the kernel's KFD driver publishes.
KFD_TOPOLOGY_NODES = "/sys/class/kfd/kfd/topology/nodes"


def amdgpu_dkms_version():
    """
    Return the amdgpu DKMS build number (e.g. 2278356 from a module version of
    "6.16.13-2278356.24.04"), or None if it cannot be determined.

    The build number is the integer immediately after the first '-' in the
    module version. dkms lists one entry per installed kernel; the entry for the
    running kernel is preferred, otherwise the highest build number is used.
    """
    try:
        result = subprocess.run(
            ["dkms", "status", "amdgpu"],
            capture_output=True,
            text=True,
            check=False,
        )
    except OSError:
        return None

    running = os.uname().release
    builds = []
    running_build = None
    for line in result.stdout.splitlines():
        # Two dkms output shapes exist: "amdgpu/<ver>, <kernel>, ..." (newer)
        # and "amdgpu, <ver>, <kernel>, ..." (older). Parse the status-free head
        # by comma, then reconcile the version/kernel split.
        head = line.split(":", 1)[0]
        parts = [p.strip() for p in head.split(",")]
        first = parts[0]
        if "/" in first:
            version = first.partition("/")[2]
            kernel = parts[1] if len(parts) > 1 else None
        else:
            version = parts[1] if len(parts) > 1 else None
            kernel = parts[2] if len(parts) > 2 else None

        if not version or "-" not in version:
            continue
        token = version.split("-", 1)[1].split(".", 1)[0]
        if not token.isdigit():
            continue

        build = int(token)
        builds.append(build)
        if kernel == running:
            running_build = build

    if running_build is not None:
        return running_build
    if builds:
        return max(builds)
    return None


def topology_reports_ais():
    """
    True if any KFD topology GPU node reports AIS P2P-DMA initialized via its
    "capability" bit (capability & KFD_CAPABILITY_AIS). Only GPU nodes are
    considered, identified by a non-zero "simd_count"; CPU nodes (simd_count 0
    or absent) are skipped.
    """
    try:
        nodes = os.listdir(KFD_TOPOLOGY_NODES)
    except OSError:
        return False

    for node in nodes:
        path = os.path.join(KFD_TOPOLOGY_NODES, node, "properties")
        try:
            with open(path, "r", encoding="utf-8") as f:
                simd_count = 0
                capability = 0
                for line in f:
                    key, _, value = line.partition(" ")
                    if key not in ("simd_count", "capability"):
                        continue
                    try:
                        parsed = int(value.strip())
                    except ValueError:
                        continue
                    if key == "simd_count":
                        simd_count = parsed
                    else:
                        capability = parsed
                if simd_count and (capability & KFD_CAPABILITY_AIS):
                    return True
        except OSError:
            continue

    return False


def config_supports_p2pdma():
    """
    Check for P2P DMA support via the kernel's CONFIG_PCI_P2PDMA config.
    """

    kr = os.uname().release
    configs_found = False

    for path, opener in [
        (f"/boot/config-{kr}", open),
        (f"/lib/modules/{kr}/build/.config", open),
        ("/proc/config.gz", gzip.open),
    ]:
        try:
            with opener(path, "rt") as f:
                configs_found = True
                for line in f:
                    if line.startswith("CONFIG_PCI_P2PDMA=y"):
                        return True
        except OSError:
            continue

    if not configs_found:
        print("No kernel config files found!", file=sys.stderr)
    return False


def kernel_supports_p2pdma():
    """
    Check for P2P DMA support.

    On amdgpu DKMS builds newer than KFD_AIS_CAPABILITY_MIN_VERSION the KFD
    topology exposes an authoritative AIS capability bit, so it alone decides the
    verdict. Older builds -- and any case where the driver version can't be
    determined -- fall back to inspecting the kernel's CONFIG_PCI_P2PDMA.
    """
    version = amdgpu_dkms_version()
    if version is not None and version > KFD_AIS_CAPABILITY_MIN_VERSION:
        return topology_reports_ais()
    return config_supports_p2pdma()


def find_hip_runtimes():
    """
    Return a mapping of HIP runtime library paths to AIS support flags
    (all initially False) by looking in the usual places.
    """

    # NOTE: CodeQL will be unhappy if you are not careful about paths
    #       in this function

    # Match both the unversioned symlink (shipped in the -dev package)
    # and the versioned runtime libraries (e.g. libamdhip64.so.6), which
    # are all that exist on runtime-only installs.
    lib_glob = "libamdhip64.so*"

    # Directories to search for HIP runtime libraries
    search_dirs = []

    # 1. Respect runtime linker paths
    for p in os.environ.get("LD_LIBRARY_PATH", "").split(":"):
        if p:
            search_dirs.append(p)

    # 2. Environment variables commonly set by ROCm or modules
    for var in ("ROCM_HOME", "ROCM_PATH", "HIP_PATH"):
        base = os.environ.get(var)
        if base:
            search_dirs += [
                os.path.join(base, "lib"),
                os.path.join(base, "lib64"),
            ]

    # 3. Standard ROCm install paths
    search_dirs += [
        "/opt/rocm/lib",
        "/opt/rocm/lib64",
    ]

    candidates = []

    # Glob each search directory for versioned and unversioned libs.
    # Clean up each directory first (removing `..`, etc. and making it
    # absolute) to keep CodeQL happy about env-derived paths. The directory
    # component is escaped so that glob metacharacters (*, ?, []) embedded in
    # an env var are matched literally rather than interpreted, and only the
    # trailing lib_glob acts as a pattern. glob.glob does not guarantee an
    # order, so each match set is sorted for deterministic probing/printing.
    for d in search_dirs:
        safe_d = os.path.abspath(os.path.normpath(d))
        if not os.path.isdir(safe_d):
            continue
        candidates += sorted(glob.glob(os.path.join(glob.escape(safe_d), lib_glob)))

    # 4. Versioned installs (/opt/rocm-5.x, etc.)
    candidates += sorted(glob.glob(f"/opt/rocm*/lib*/{lib_glob}"))

    # 5. ROCm >= 7.11 versioned core installs (/opt/rocm/core-X.Y/lib),
    #    where /opt/rocm is a real directory rather than a symlink.
    candidates += sorted(glob.glob(f"/opt/rocm*/core-*/lib*/{lib_glob}"))

    # 6. Fall back to the ldconfig cache, which catches installs in
    #    non-standard prefixes registered via /etc/ld.so.conf.d
    cached = ctypes.util.find_library("amdhip64")
    if cached:
        candidates.append(cached)

    # Resolve, dedup, and drop paths that don't exist.
    #
    # Absolute paths are canonicalized with realpath so that symlinks
    # (the unversioned .so pointing at a versioned one, or several rocm
    # directories pointing at a single install) collapse to one entry.
    # find_library may return a bare soname with no path to check, but
    # CDLL can still load it by name, so it is kept as-is.
    resolved = []
    for path in candidates:
        if os.path.isabs(path):
            real = os.path.realpath(path)
            if os.path.exists(real):
                resolved.append(real)
        else:
            resolved.append(path)

    # dict.fromkeys dedups while preserving insertion order.
    return dict.fromkeys(resolved, False)


def hip_runtime_supports_ais():
    """
    Check if hipAmdFileRead and hipAmdFileWrite are available in HIP.

    Returns the mapping of discovered HIP library paths to a bool
    indicating whether that library supports AIS.
    """
    hip_libraries = find_hip_runtimes()

    hipError_t = ctypes.c_int
    hipDriverProcAddressQueryResult = ctypes.c_int

    # Check for AIS functions in the list of found HIP libraries
    for hip_path in hip_libraries:

        try:
            hip = ctypes.CDLL(hip_path)

            # Resolving these symbols can raise AttributeError on a HIP runtime
            # too old to export them (e.g. hipGetProcAddress). Treat such a
            # library as simply not AIS-capable rather than crashing.
            hip.hipRuntimeGetVersion.argtypes = [ctypes.POINTER(ctypes.c_int)]
            hip.hipRuntimeGetVersion.restype = hipError_t

            hip.hipGetProcAddress.argtypes = [
                ctypes.c_char_p,
                ctypes.POINTER(ctypes.c_void_p),
                ctypes.c_int,
                ctypes.c_uint64,
                ctypes.POINTER(hipDriverProcAddressQueryResult),
            ]
            hip.hipGetProcAddress.restype = hipError_t

            hip.hipGetErrorString.argtypes = [hipError_t]
            hip.hipGetErrorString.restype = ctypes.c_char_p
        except (OSError, AttributeError):
            continue

        version = ctypes.c_int()
        err = hip.hipRuntimeGetVersion(ctypes.byref(version))
        if err != 0:
            err_bytes = hip.hipGetErrorString(err)
            err_str = err_bytes.decode("utf-8") if err_bytes else "unknown error"
            print(
                f"hipRuntimeGetVersion failed for {hip_path} "
                f"with err code {err} ({err_str})",
                file=sys.stderr,
            )
            continue

        # Track whether all required AIS symbols are available in this library
        supported = True

        for symbol in [b"hipAmdFileWrite", b"hipAmdFileRead"]:
            func_ptr = ctypes.c_void_p()
            symbol_status = hipDriverProcAddressQueryResult()
            err = hip.hipGetProcAddress(
                symbol,
                ctypes.byref(func_ptr),
                version.value,
                0,
                ctypes.byref(symbol_status),
            )
            # A symbol is only present if the call succeeded, the query result
            # is HIP_GET_PROC_ADDRESS_SUCCESS (0), and a non-null pointer came
            # back. Checking all three avoids false positives across HIP
            # implementations.
            if err != 0 or symbol_status.value != 0 or not func_ptr.value:
                supported = False
                break

        if supported:
            hip_libraries[hip_path] = True

    return hip_libraries


def amdgpu_supports_ais():
    """
    Check if kfd_ais_rw_file is in the kernel's symbol table
    """
    try:
        with open("/proc/kallsyms", "r", encoding="utf-8") as kallsyms:
            for line in kallsyms:
                if "kfd_ais_rw_file" in line:
                    return True
    except FileNotFoundError:
        print(
            "Unable to open /proc/kallsyms: file not found",
            file=sys.stderr,
        )
    except PermissionError:
        print(
            "Unable to open /proc/kallsyms: permission denied",
            file=sys.stderr,
        )
    return False


class Mount:  # pylint: disable=too-few-public-methods
    """A single parsed /proc/self/mountinfo entry."""

    def __init__(self, devno, mountpoint, fstype, source, options):
        self.devno = devno  # "maj:min"
        self.mountpoint = mountpoint
        self.fstype = fstype
        self.source = source
        self.options = options  # combined mount + super options, comma-joined

    def option(self, name):
        """Return the value of mount option `name`, or None if unset/valueless."""
        for opt in self.options.split(","):
            if opt == name:
                return ""
            if opt.startswith(name + "="):
                return opt.split("=", 1)[1]
        return None


def unescape(field):
    """Decode the octal escapes (\\040 etc.) mountinfo uses for spaces/tabs."""
    out = []
    i = 0
    while i < len(field):
        if field[i] == "\\" and field[i + 1 : i + 4].isdigit():
            out.append(chr(int(field[i + 1 : i + 4], 8)))
            i += 4
        else:
            out.append(field[i])
            i += 1
    return "".join(out)


def parse_mountinfo():
    """Parse /proc/self/mountinfo into a list of Mount objects."""
    mounts = []
    with open("/proc/self/mountinfo", "r", encoding="utf-8") as f:
        for line in f:
            fields = line.split()
            # Fields up to the "-" separator are fixed; optional fields vary.
            try:
                sep = fields.index("-")
            except ValueError:
                continue
            devno = fields[2]
            mountpoint = unescape(fields[4])
            mount_opts = fields[5]
            fstype = fields[sep + 1]
            source = unescape(fields[sep + 2])
            super_opts = fields[sep + 3] if len(fields) > sep + 3 else ""
            options = ",".join(o for o in (mount_opts, super_opts) if o)
            mounts.append(Mount(devno, mountpoint, fstype, source, options))
    return mounts


def fs_supported(mount):
    """
    Mirror mountinfo.cpp: True if the filesystem type/journaling qualifies.

    ext4 qualifies only with data=ordered (the default when no data= option is
    present); xfs always qualifies.
    """
    if mount.fstype == "xfs":
        return True
    if mount.fstype == "ext4":
        data = mount.option("data")
        # A missing data= is treated as ordered to match libmount/mountinfo.cpp,
        # which defaults to ordered when the option is absent. The effective mode
        # also lives in /proc/fs/ext4/<dev>/options, but we deliberately mirror
        # the C reference so our verdict matches what hipFile decides at runtime.
        return data is None or data == "ordered"
    return False


def fstype_label(mount):
    """
    Render the filesystem type, folding ext4's journal mode in parentheses.

    ext4 always journals, so we surface the data= mode (defaulting to the
    "ordered" ext4 default) as e.g. "ext4 (ordered)". Other filesystems are
    shown bare since the journal mode either doesn't apply or isn't expressed
    as a data= option.
    """
    if mount.fstype == "ext4":
        return f"ext4 ({mount.option('data') or 'ordered'})"
    return mount.fstype


def probe_odirect(mountpoint):
    """
    Determine O_DIRECT support by opening a temp file with O_DIRECT.

    Returns True (supported), False (rejected with EINVAL), or None when it
    could not be verified (e.g. read-only or no write permission).
    """
    o_direct = getattr(os, "O_DIRECT", 0)
    if o_direct == 0:
        # os.O_DIRECT is only defined when the C library Python was built
        # against defines O_DIRECT; its absence reflects that build, not
        # whether the kernel/FS supports O_DIRECT. Without the flag we can't
        # construct the test open, so we can't verify support either way.
        return None

    flags = os.O_RDWR | os.O_CREAT | os.O_EXCL | os.O_CLOEXEC | o_direct
    path = os.path.join(mountpoint, f".ais-odirect-{os.getpid()}-{os.urandom(4).hex()}")
    fd = None
    try:
        fd = os.open(path, flags, 0o600)
        return True
    except OSError as e:
        if e.errno == errno.EINVAL:
            return False
        # EROFS / EACCES / EPERM / ENOSPC etc. -> can't tell
        return None
    finally:
        if fd is not None:
            os.close(fd)
        try:
            os.unlink(path)
        except OSError:
            pass


def backing_storage(devno):
    """
    Resolve a maj:min to (backing_type, disk_name) via sysfs.

    `backing_type` is a single short label describing what the filesystem
    resides on: device-mapper targets are classified from their dm/uuid
    (lvm, mpath, crypt, dm), software RAID as md, and raw block devices from
    their kernel name (nvme, virtio, scsi, loop). Returns (None, None) when
    sysfs cannot be read.
    """
    try:
        real = os.path.realpath(f"/sys/dev/block/{devno}")
    except OSError:
        return None, None
    name = os.path.basename(real)

    # device-mapper targets (lvm, multipath, crypt, ...) carry a dm/uuid whose
    # prefix names the target type.
    try:
        with open(os.path.join(real, "dm", "uuid"), "r", encoding="utf-8") as f:
            uuid = f.read().strip()
    except OSError:
        uuid = None
    if uuid is not None:
        prefix = uuid.split("-", 1)[0].lower()
        backing = {"lvm": "lvm", "mpath": "mpath", "crypt": "crypt"}.get(prefix, "dm")
    elif name.startswith("md"):
        backing = "md"
    elif name.startswith("nvme"):
        backing = "nvme"
    elif name.startswith(("vd", "xvd")):
        backing = "virtio"
    elif name.startswith(("sd", "hd", "sr")):
        backing = "scsi"
    elif name.startswith("loop"):
        backing = "loop"
    else:
        backing = name or None

    return backing, name


def backing_supported(backing):
    """
    Whether the backing storage allows hipFile's direct P2P-DMA path.

    hipFile is limited to local NVMe (INSTALL.md), and even multipath NVMe is
    unsupported. Any interposing layer (LVM/dm, multipath, dm-crypt, MD RAID,
    loopback) breaks the direct path to the device, so the kernel rejects the IO
    at runtime (pcie_p2pdma_distance < 0) even when the filesystem and O_DIRECT
    checks pass. Returns None when the backing could not be determined.
    """
    if backing is None:
        return None
    return backing == "nvme"


def tri(value, unknown="?"):
    """Render an Optional[bool] as a short cell."""
    if value is True:
        return "yes"
    if value is False:
        return "no"
    return unknown


def collect(mounts):
    """Build a list of result rows for the block-backed mounts."""
    rows = []
    for m in mounts:
        if m.fstype in PSEUDO_FSTYPES:
            continue
        # Only consider local block devices. Network filesystems (NFS
        # "server:/export", CIFS "//server/share", sshfs, ...) are not
        # block-backed.
        if not m.source.startswith("/dev/"):
            continue

        backing, disk = backing_storage(m.devno)
        supported_fs = fs_supported(m)
        backing_ok = backing_supported(backing)

        # Only probe fs types that can qualify; unsupported fs is never capable.
        odirect = probe_odirect(m.mountpoint) if supported_fs else None

        # A volume is capable only when the filesystem qualifies, the backing is
        # a direct local NVMe (interposing layers like LVM disqualify it), and
        # O_DIRECT works. Unknown backing leaves the verdict unverified.
        if not supported_fs or backing_ok is False:
            capable = False
        elif backing_ok is None:
            capable = None
        else:
            capable = odirect

        rows.append(
            {
                "mountpoint": m.mountpoint,
                "fstype": fstype_label(m),
                "device": disk or m.source,
                "backing": backing,
                "odirect": odirect,
                "capable": capable,
            }
        )
    return rows


def print_volume_table(rows):
    """Print the collected volume rows as an aligned table."""
    headers = ["MOUNTPOINT", "FSTYPE", "DEVICE", "BACKING", "O_DIRECT", "HIPFILE"]
    table = [headers]
    for r in rows:
        table.append(
            [
                r["mountpoint"],
                r["fstype"],
                r["device"],
                r["backing"] or "?",
                tri(r["odirect"], unknown="unverified"),
                tri(r["capable"], unknown="unverified"),
            ]
        )

    widths = [max(len(row[i]) for row in table) for i in range(len(headers))]
    for row in table:
        print("  ".join(cell.ljust(widths[i]) for i, cell in enumerate(row)))


def capable_volumes():
    """
    Enumerate block-backed mounts and score them for hipFile capability.

    Returns (rows, any_capable) where rows is the list of per-volume result
    dicts and any_capable is True if at least one volume is confidently capable.
    """
    rows = collect(parse_mountinfo())
    return rows, any(r["capable"] is True for r in rows)


def main():
    """
    Parse command-line arguments, check AIS support in kernel/libraries and
    mounted volumes, optionally print the results, and return an exit code
    indicating whether all required components support AIS.
    """

    parser = argparse.ArgumentParser(
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=(
            "O_DIRECT/HIPFILE columns are probed by opening a temp file with "
            "O_DIRECT on each volume.\nThis needs write access to the mountpoint "
            "(typically root); without it the result is\nshown as 'unverified' "
            "rather than supported/unsupported. Re-run as root to verify."
        ),
    )
    output_group = parser.add_mutually_exclusive_group()
    output_group.add_argument(
        "-q", "--quiet", action="store_true", help="Silence regular output"
    )
    output_group.add_argument(
        "-v",
        "--verbose",
        action="store_true",
        help="Also list the discovered HIP libraries",
    )
    args = parser.parse_args()

    hip_libraries = hip_runtime_supports_ais()
    volume_rows, volumes_ok = capable_volumes()

    component_support = [
        ("Kernel P2PDMA support", kernel_supports_p2pdma()),
        ("HIP runtime", any(hip_libraries.values())),
        ("amdgpu", amdgpu_supports_ais()),
        ("hipFile-capable volume", volumes_ok),
    ]

    if not args.quiet:
        u = os.uname()
        print()
        print(u.sysname, u.nodename, u.release, u.version, u.machine)

        if args.verbose:
            print()
            print("Found these HIP libraries (some may refer to the same library):")
            for lib, support in hip_libraries.items():
                if support:
                    pretty_supported = "supported"
                else:
                    pretty_supported = "NOT supported"
                print(f"\t{lib} (AIS {pretty_supported})")

        print()
        print("Mounted volumes:")
        if volume_rows:
            print_volume_table(volume_rows)
        else:
            print("No block-backed volumes found.")

        print()
        print("AIS support in:")
        for name, supported in component_support:
            print(f"\t{name:<24}: {supported}")

    return int(not all(supported for (_, supported) in component_support))


if __name__ == "__main__":
    sys.exit(main())
