"""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 platform


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, soname: str, dllname: str):
        self.shortname = shortname
        self.package = ALL_PACKAGES[package_name]
        self.posix_relpath = "lib"
        self.windows_relpath = "bin"
        self.soname = soname
        self.dllname = dllname
        assert shortname not in ALL_LIBRARIES
        ALL_LIBRARIES[shortname] = self

    def __repr__(self):
        return f"{self.shortname}(soname={self.soname}, dllname={self.dllname}, 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


# 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 discover_current_target_family() -> str | None:
    # TODO: Implement dynamic discovery.
    return None


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): Use patterns for version suffixes and platform differences.

# Public libraries.
LibraryEntry("amdhip64", "core", "libamdhip64.so.6", "amdhip64_6.dll")
LibraryEntry("hiprtc", "core", "libhiprtc.so.6", "hiprtc0605.dll")
LibraryEntry("roctx64", "core", "libroctx64.so.4", "")
LibraryEntry("rocprofiler-sdk-roctx", "core", "librocprofiler-sdk-roctx.so.1", "")

LibraryEntry("amd_comgr", "core", "libamd_comgr.so.3", "amd_comgr0605.dll")
LibraryEntry("hipblas", "libraries", "libhipblas.so.3", "libhipblas.dll")
LibraryEntry("hipfft", "libraries", "libhipfft.so.0", "hipfft.dll")
LibraryEntry("hiprand", "libraries", "libhiprand.so.1", "hiprand.dll")
LibraryEntry("hipsparse", "libraries", "libhipsparse.so.4", "hipsparse.dll")
LibraryEntry("hipsolver", "libraries", "libhipsolver.so.1", "hipsolver.dll")
LibraryEntry("rccl", "libraries", "librccl.so.1", "")
LibraryEntry("hipblaslt", "libraries", "libhipblaslt.so.1", "hipblaslt.dll")
LibraryEntry("miopen", "libraries", "libMIOpen.so.1", "MIOpen.dll")

# 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.0.0.dev0+e96d36b9b628476463ef6cecee752f601052a4d2'
PY_PACKAGE_SUFFIX_NONCE = ''
DEFAULT_TARGET_FAMILY = 'gfx110X-dgpu'
AVAILABLE_TARGET_FAMILIES.append('gfx110X-dgpu')
THIS_TARGET_FAMILY = None
THIS_PACKAGE_ENTRY = ALL_PACKAGES['meta']
