# Copyright Advanced Micro Devices, Inc.
# SPDX-License-Identifier: MIT

"""Test of the library trees."""

"""Installation package tests for the core package."""

import importlib
import os
from pathlib import Path
import platform
import subprocess
import sys
import unittest

from .. import _dist_info as di
from . import utils

import rocm_sdk


class ROCmDevelTest(unittest.TestCase):
    def testInstallationLayout(self):
        """The `rocm_sdk` and devel module must be siblings on disk."""
        sdk_path = Path(rocm_sdk.__file__)
        self.assertEqual(
            sdk_path.name,
            "__init__.py",
            msg="Expected `rocm_sdk` module to be a non-namespace package",
        )
        import rocm_sdk_devel

        devel_path = Path(rocm_sdk_devel.__file__)
        self.assertEqual(
            devel_path.name,
            "__init__.py",
            msg=f"Expected `rocm_sdk_devel` module to be a non-namespace package",
        )
        self.assertEqual(
            sdk_path.parent.parent,
            devel_path.parent.parent,
            msg="Paths are not siblings",
        )

    def testCLIPathBin(self):
        cmd = [sys.executable, "-m", "rocm_sdk", "path", "--bin"]
        output = utils.run_command(cmd, capture=True).decode().strip()
        path = Path(output)
        self.assertTrue(path.exists(), msg=f"Expected bin path {path} to exist")

    def testCLIPathCMake(self):
        cmd = [sys.executable, "-m", "rocm_sdk", "path", "--cmake"]
        output = utils.run_command(cmd, capture=True).decode().strip()
        path = Path(output)
        self.assertTrue(path.exists(), msg=f"Expected cmake path {path} to exist")
        hip_file = path / "hip" / "hip-config.cmake"
        self.assertTrue(
            hip_file.exists(), msg=f"Expected hip config to exist {hip_file}"
        )

    def testCLIPathRoot(self):
        cmd = [sys.executable, "-m", "rocm_sdk", "path", "--root"]
        output = utils.run_command(cmd, capture=True).decode().strip()
        path = Path(output)
        self.assertTrue(path.exists(), msg=f"Expected root path {path} to exist")
        bin_path = path / "bin"
        self.assertTrue(bin_path.exists(), msg=f"Expected bin path {bin_path} to exist")

    def testCLIUsesDevelRootPath(self):
        root_path_output = (
            utils.run_command(
                [sys.executable, "-m", "rocm_sdk", "path", "--root"], capture=True
            )
            .decode()
            .strip()
        )
        root_path = Path(root_path_output)

        # CLI scripts by default run from _rocm_sdk_core.
        # When the devel package is installed they should run from _rocm_sdk_devel.
        rocmpath_output = (
            utils.run_command(["hipconfig", "--rocmpath"], capture=True)
            .decode()
            .strip()
        )
        rocmpath = Path(rocmpath_output)
        self.assertEqual(
            root_path,
            rocmpath,
            msg=f"Expected `hipconfig --rocmpath` to return {root_path}, not {rocmpath}",
        )

    @unittest.skipIf(
        platform.system() == "Windows", "root LLVM symlink only exists on Linux"
    )
    def testRootLLVMSymlinkExists(self):
        # We had a bug where the root llvm/ symlink, which is for backwards compat,
        # was not materialized. Verify it is.
        cmd = [sys.executable, "-m", "rocm_sdk", "path", "--root"]
        output = utils.run_command(cmd, capture=True).decode().strip()
        path = Path(output) / "llvm" / "bin" / "clang++"
        self.assertTrue(path.exists(), msg=f"Expected {path} to exist")

    def testSharedLibrariesLoad(self):
        # Make sure the devel package is expanded.
        cmd = [sys.executable, "-m", "rocm_sdk", "path", "--root"]
        _ = utils.run_command(cmd, capture=True).decode().strip()

        # Ensure that the platform package exists now.
        mod_name = di.ALL_PACKAGES["devel"].get_py_package_name(
            target_family=di.determine_target_family()
        )
        mod = importlib.import_module(mod_name)
        utils.assert_is_physical_package(mod)
        so_paths = utils.get_module_shared_libraries(mod)

        self.assertTrue(
            so_paths, msg="Expected core package to contain shared libraries"
        )

        for so_path in so_paths:
            if "amd_smi" in str(so_path) or "goamdsmi" in str(so_path):
                # TODO: Library preloads for amdsmi need to be implement.
                # Though this is not needed for the amd-smi client.
                continue
            if "clang_rt" in str(so_path):
                # clang_rt and sanitizer libraries are not all intended to be
                # loadable arbitrarily.
                continue
            if "libhipsolver_fortran" in str(so_path):
                # Currently fails to load unless libgfortran.so.5 exists on the system.
                # TODO(#3115): Decide if this test should be permanently
                #     disabled or fixed and then re-enabled somehow. This
                #     library may only be used by tests and we might not care
                #     about it failing to load standalone.
                continue
            if "libLLVMOffload" in str(so_path):
                # recent addition from upstream, issue tracked in
                # https://github.com/ROCm/TheRock/issues/2537
                continue
            if "lib/roctracer" in str(so_path) or "share/roctracer" in str(so_path):
                # Internal roctracer libraries are meant to be pre-loaded
                # explicitly and cannot necessarily be loaded standalone.
                continue
            if (
                "lib/rocprofiler-sdk/" in str(so_path)
                or "libexec/rocprofiler-sdk/" in str(so_path)
                or "libpyrocpd" in str(so_path)
                or "libpyroctx" in str(so_path)
            ):
                # Internal rocprofiler-sdk libraries are meant to be pre-loaded
                # explicitly and cannot necessarily be loaded standalone.
                continue
            if "libtest_linking_lib" in str(so_path):
                # rocprim unit tests, not actual library files
                continue
            if "opencl" in str(so_path):
                # We use OpenCL ICD from distro rather than TheRock
                # and we do not build it
                continue

            extra_setup = ""
            if (
                "hipdnn_plugins" in str(so_path) or "test_plugins" in str(so_path)
            ) and platform.system() == "Windows":
                # hipdnn plugins have dependencies on other libraries (e.g. miopen).
                # In a real-world scenario, hipdnn_backend loads these plugins, and
                # the dependencies are found because they reside in the same directory
                # (or are otherwise resolvable).
                # To simulate this loading behavior in the test:
                # - On Linux, RPATH ($ORIGIN/../../) handles dependency resolution.
                # - On Windows, we must manually add the library directory (calculated
                #   relative to the plugin) via add_dll_directory, as there is no RPATH equivalent.
                # We assume the plugin is at .../{lib|bin}/hipdnn_plugins/engines/plugin.so
                # and the dependencies are at .../{lib|bin}.
                lib_dir = str(so_path.parents[2]).replace("\\", "\\\\")
                extra_setup = f"import os; os.add_dll_directory('{lib_dir}') if hasattr(os, 'add_dll_directory') else None; "

            with self.subTest(msg="Check shared library loads", so_path=so_path):
                # Load each in an isolated process because not all libraries in the tree
                # are designed to load into the same process (i.e. LLVM runtime libs,
                # etc).
                command = (
                    extra_setup + "import ctypes; import sys; ctypes.CDLL(sys.argv[1])"
                )

                subprocess.check_call([sys.executable, "-c", command, str(so_path)])
