"""Distribution information.

This file is typically auto-generated by the build system, but for the purposes
of bootstrapping, we are including it inline for the moment.
"""

import importlib.util
import os
import subprocess
from pathlib import Path


CACHED_TARGET_FAMILY: str | None = None


class LibraryEntry:
    """Defines a public library that can be located by name within the overall
    distribution."""

    def __init__(
        self,
        shortname: str,
        package_name: str,
        so_pattern: str,
        dll_pattern: str,
        posix_relpath="lib",
    ):
        self.shortname = shortname
        self.package = ALL_PACKAGES[package_name]
        self.posix_relpath = posix_relpath
        self.windows_relpath = "bin"
        self.so_pattern = so_pattern
        self.dll_pattern = dll_pattern
        assert shortname not in ALL_LIBRARIES
        ALL_LIBRARIES[shortname] = self

    def __repr__(self):
        return f"{self.shortname}(so_pattern={self.so_pattern}, dll_pattern={self.dll_pattern}, package={self.package})"


class PackageEntry:
    """Defines a package known to the SDK.

    Args:
      logical_name: Short name referring to this package within the SDK.
      dist_package_name: Templated dist package name, including '{target_family}' if
        GPU target specific.
      pure_py_package_name: Name of the pure-python API package. All GPU target specific
        dist packages will export the same pure-python package. Since they are
        identical, first wins.
      template_directory: Directory name in the build system that provides the template
        for this package.
      required: Whether this package is required.
    """

    def __init__(
        self,
        logical_name: str,
        dist_package_template: str,
        *,
        pure_py_package_name: str,
        template_directory: str,
        required: bool = False,
    ):
        self.logical_name = logical_name
        self.dist_package_template = dist_package_template
        self.pure_py_package_name = pure_py_package_name
        self.template_directory = template_directory
        self.required = required
        if logical_name in ALL_PACKAGES:
            raise ValueError(f"Package already defined: {logical_name}")
        ALL_PACKAGES[logical_name] = self

    @property
    def is_target_specific(self) -> bool:
        return "{target_family}" in self.dist_package_template

    def get_dist_package_name(self, target_family: str | None = None) -> str:
        if self.is_target_specific and target_family is None:
            raise ValueError(
                f"Package {self.logical_name} is target specific, but no target specified"
            )
        kwargs = {}
        if target_family is not None:
            kwargs["target_family"] = target_family
        return self.dist_package_template.format(**kwargs)

    def get_dist_package_require(self, target_family: str | None = None) -> str:
        return self.get_dist_package_name(target_family) + f"=={__version__}"

    def get_py_package_name(self, target_family: str | None = None) -> str:
        dist_name = self.get_dist_package_name(target_family)
        return "_" + dist_name.replace("-", "_") + PY_PACKAGE_SUFFIX_NONCE

    def get_py_package(self, target_family: str | None = None):
        return importlib.util.find_spec(self.get_py_package_name(target_family))

    def has_py_package(self, target_family: str | None = None) -> bool:
        return self.get_py_package(target_family) is not None

    def __repr__(self):
        return self.dist_package_template


def discover_current_target_family() -> str | None:
    """Attempts to query the current target family via the 'offload-arch' tool."""

    try:
        import sysconfig

        # offload-arch is expected to be installed in the Python 'scripts'
        # directory, which will vary depending on the platform and whether or
        # not a virtual environment is used, for example:
        #   Linux system:   /usr/local/bin
        #   Linux venv:     .venv/bin
        #   Windows system: C:\Users\...\Python313\Scripts
        #   Windows venv:   .venv\Scripts
        # It might also be provided by an install of LLVM (e.g. as part of
        # Visual Studio on Windows), so prepend the scripts dir to PATH.
        scripts_path = Path(sysconfig.get_path("scripts"))
        env = os.environ
        env["PATH"] = str(scripts_path) + os.path.pathsep + env.get("PATH", "")
        result = subprocess.check_output(["offload-arch"], env=env, text=True)

        if result:
            arch_set = set(result.strip().split("\n"))
            suffixes = ["-all", "-dgpu", "-igpu", "-dcgpu"]
            for arch in arch_set:
                # There may be multiple architecture supported on the system.
                # This will select the first matching family.
                arch_family = arch[:-1] + "X"
                for suffix in suffixes:
                    target_family = arch_family + suffix
                    if target_family in AVAILABLE_TARGET_FAMILIES:
                        return target_family
                if arch in AVAILABLE_TARGET_FAMILIES:
                    return arch
    except subprocess.CalledProcessError as e:
        print(f"[WARNING] offload-arch failed with return code {e.returncode}")
        print(f"[stderr] {e.output}")
    except FileNotFoundError:
        print(f"[WARNING] failed to run offload-arch: binary not found.")
    except Exception as e:
        print(f"[WARNING] Unexpected error running offload-arch: {e}")
    return None


# Resolve the build target family. This consults a list of things in increasing
# order of specificity:
#   1. "ROCM_SDK_TARGET_FAMILY" environment variable
#   2. Dynamically discovered/most salient target family on the actual system
#   3. dist_info.DEFAULT_TARGET_FAMILY
def determine_target_family() -> str:
    global CACHED_TARGET_FAMILY
    if CACHED_TARGET_FAMILY is not None:
        return CACHED_TARGET_FAMILY
    target_family = os.getenv("ROCM_SDK_TARGET_FAMILY")
    if target_family is None:
        target_family = discover_current_target_family()
        if target_family is None:
            target_family = DEFAULT_TARGET_FAMILY
    assert target_family is not None
    if target_family not in AVAILABLE_TARGET_FAMILIES:
        raise ValueError(
            f"Requested ROCM_SDK_TARGET_FAMILY={target_family} is "
            f"not available in the distribution (available: "
            f"{', '.join(AVAILABLE_TARGET_FAMILIES)})"
        )
    CACHED_TARGET_FAMILY = target_family
    return target_family


# All packages that are part of the distribution.
ALL_PACKAGES: dict[str, PackageEntry] = {}
ALL_LIBRARIES: dict[str, LibraryEntry] = {}

# Always available packages.
PackageEntry(
    "meta",
    "rocm",
    pure_py_package_name="rocm_sdk",
    template_directory="rocm",
    required=True,
)
PackageEntry(
    "core",
    "rocm-sdk-core",
    pure_py_package_name="rocm_sdk_core",
    template_directory="rocm-sdk-core",
    required=True,
)
PackageEntry(
    "libraries",
    "rocm-sdk-libraries-{target_family}",
    pure_py_package_name="rocm_sdk_libraries",
    template_directory="rocm-sdk-libraries",
    required=False,
)
PackageEntry(
    "devel",
    "rocm-sdk-devel",
    pure_py_package_name="rocm_sdk_devel",
    template_directory="rocm-sdk-devel",
    required=False,
)

# TODO(#703,#1057): Use patterns for version suffixes and platform differences too?

# Public libraries.
LibraryEntry("amdhip64", "core", "libamdhip64.so*", "amdhip64*.dll")
# The DLL glob here uses '0' from the version to avoid matching 'hiprtc-builtins'.
# If DLLs with no version suffix are later added we will need a different pattern.
LibraryEntry("hiprtc", "core", "libhiprtc.so*", "hiprtc0*.dll")
LibraryEntry("roctx64", "core", "libroctx64.so*", "")
LibraryEntry("rocprofiler-sdk-roctx", "core", "librocprofiler-sdk-roctx.so*", "")
LibraryEntry("roctracer64", "core", "libroctracer64.so*", "")
LibraryEntry(
    "rocm_sysdeps_liblzma",
    "core",
    "librocm_sysdeps_liblzma.so.*",
    "",
    "lib/rocm_sysdeps/lib",
)
LibraryEntry(
    "rocm-openblas",
    "core",
    "librocm-openblas.so.*",
    "rocm-openblas*.dll",
    "lib/host-math/lib",
)
LibraryEntry("amd_comgr", "core", "libamd_comgr.so*", "amd_comgr*.dll")
LibraryEntry("hipblas", "libraries", "libhipblas.so*", "*hipblas*.dll")
LibraryEntry("hipblaslt", "libraries", "libhipblaslt.so*", "*hipblaslt*.dll")
LibraryEntry("hipfft", "libraries", "libhipfft.so*", "hipfft*.dll")
LibraryEntry("hiprand", "libraries", "libhiprand.so*", "hiprand*.dll")
LibraryEntry("hipsparse", "libraries", "libhipsparse.so*", "hipsparse*.dll")
LibraryEntry("hipsparselt", "libraries", "libhipsparselt.so*", "")
LibraryEntry("hipsolver", "libraries", "libhipsolver.so*", "hipsolver*.dll")
LibraryEntry("rccl", "libraries", "librccl.so*", "")
LibraryEntry("miopen", "libraries", "libMIOpen.so*", "MIOpen*.dll")

# Others we may want:
# hiprtc-builtins
# rocblas
# rocfft
# rocrand
# rocsolver
# rocsparse

# Overall ROCM package version.
__version__ = "DEFAULT"

# Nonce added to the backend packages which encodes the version. This is
# typically empty for development distributions. Only backend packages with
# a matching nonce will be considered for use by this meta package.
PY_PACKAGE_SUFFIX_NONCE: str = "_DEFAULT"

# If a target family cannot be found or is not relevant (i.e. building devel
# packages on a gpu-less system), this is the default target family.
DEFAULT_TARGET_FAMILY: str = "DEFAULT"

# All available target families that this distribution has available.
AVAILABLE_TARGET_FAMILIES: list[str] = []
__version__ = '7.11.0.dev0+027a894bb27ec70d15c996e05bc5db0ee9026616'
PY_PACKAGE_SUFFIX_NONCE = ''
DEFAULT_TARGET_FAMILY = 'gfx1151'
AVAILABLE_TARGET_FAMILIES.append('gfx1151')
THIS_TARGET_FAMILY = None
THIS_PACKAGE_ENTRY = ALL_PACKAGES['meta']
