import sys
from typing import Any, Dict, List

import cupy_builder.install_build as build
import cupy_builder.install_utils as utils
from cupy_builder import Context


class Feature:
    NOT_AVAILABLE = -1
    _UNDETERMINED = -100

    def __init__(self, ctx: Context):
        # Name of the feature.
        self.name = ''

        # When True, fail the build if the feature is unavailable.
        self.required = False

        # List of Cython modules.
        self.modules: List[str] = []

        # C/C++ headers required for the feature.
        # This is used only for testing availability of the feature.
        self.includes: List[str] = []

        # Libraries (shared/static) on the search path to be linked.
        self.libraries: List[str] = []

        # Static libraries (manually searched) to be linked.
        self.static_libraries: List[str] = []

        # Version of the feature.
        self._version: Any = self._UNDETERMINED

    def configure(self, compiler: Any, settings: Any) -> bool:
        # Fill `self._version` with the version or NOT_AVAILABLE
        self._version = None
        return True

    def get_version(self) -> Any:
        assert self._version != self._UNDETERMINED, 'not configured yet'
        return self._version

    def __contains__(self, key: Any) -> bool:
        # TODO(kmaehashi): Remove this transient function.
        if not isinstance(key, str):
            return False
        try:
            self.__getitem__(key)
        except AttributeError:
            return False
        return True

    def __getitem__(self, key: str) -> Any:
        # TODO(kmaehashi): Remove this transient function.
        if key == 'file':
            return self.modules
        elif key == 'include':
            return self.includes
        return getattr(self, key)


def _from_dict(d: Dict[str, Any], ctx: Context) -> Feature:
    # Define a feature from dict.
    # TODO(kmaehashi): Remove this transient function.
    f = Feature(ctx)
    f.name = d['name']
    f.required = d.get('required', False)
    f.libraries = d['libraries']
    f.static_libraries = d.get('static_libraries', [])

    # Note: the following are renamed
    f.modules = d['file']
    f.includes = d['include']
    if 'check_method' in d:
        f.configure = d['check_method']  # type: ignore
        f._version = None
        if 'version_method' in d:
            f.get_version = d['version_method']  # type: ignore
    return f


# The value of the key 'file' is a list that contains extension names
# or tuples of an extension name and a list of other sources files
# required to build the extension such as .cpp files and .cu files.
#
#   <extension name> | (<extension name>, a list of <other source>)
#
# The extension name is also interpreted as the name of the Cython
# source file required to build the extension with appending '.pyx'
# file extension.

_cuda_files = [
    'cupy_backends.cuda.api.driver',
    'cupy_backends.cuda.api._driver_enum',
    'cupy_backends.cuda.api.runtime',
    'cupy_backends.cuda.api._runtime_enum',
    'cupy_backends.cuda.libs.cublas',
    'cupy_backends.cuda.libs.curand',
    'cupy_backends.cuda.libs.cusparse',
    'cupy_backends.cuda.libs.nvrtc',
    'cupy_backends.cuda.stream',
    'cupy_backends.cuda._softlink',
    'cupy._core._accelerator',
    'cupy._core._carray',
    'cupy._core._cub_reduction',
    'cupy._core._dtype',
    'cupy._core._fusion_kernel',
    'cupy._core._fusion_thread_local',
    'cupy._core._fusion_trace',
    'cupy._core._fusion_variable',
    'cupy._core._kernel',
    'cupy._core._memory_range',
    'cupy._core._optimize_config',
    'cupy._core._reduction',
    'cupy._core._routines_binary',
    'cupy._core._routines_indexing',
    'cupy._core._routines_linalg',
    'cupy._core._routines_logic',
    'cupy._core._routines_manipulation',
    'cupy._core._routines_math',
    'cupy._core._routines_sorting',
    'cupy._core._routines_statistics',
    'cupy._core._scalar',
    'cupy._core.core',
    'cupy._core.flags',
    'cupy._core.internal',
    'cupy._core.fusion',
    'cupy._core.new_fusion',
    'cupy._core.numpy_allocator',
    'cupy._core.raw',
    'cupy.cuda.common',
    'cupy.cuda.cufft',
    'cupy.cuda.device',
    'cupy.cuda.memory',
    'cupy.cuda.memory_hook',
    'cupy.cuda.pinned_memory',
    'cupy.cuda.function',
    'cupy.cuda.stream',
    'cupy.cuda.graph',
    'cupy.cuda.texture',
    'cupy.fft._cache',
    'cupy.fft._callback',
    'cupy.lib._polynomial',
    'cupy._util',
]

# Libraries required for cudart_static
_cudart_static_libs = (
    (['pthread', 'rt', 'dl'] if sys.platform == 'linux' else [])
)


def get_features(ctx: Context) -> Dict[str, Feature]:
    # We handle nvtx (and likely any other future support) here, because
    # the HIP stubs (hip/cupy_*.h) would cause many symbols
    # to leak into all these modules even if unused. It's easier for all of
    # them to link to the same set of shared libraries.
    HIP_cuda_nvtx_cusolver = {
        # TODO(leofang): call this "rocm" or "hip" to avoid confusion?
        'name': 'cuda',
        'required': True,
        'file': _cuda_files + [
            'cupy_backends.cuda.libs.nvtx',
            'cupy_backends.cuda.libs.cusolver',
            'cupyx.cusolver',
        ],
        'include': [
            'hip/hip_runtime_api.h',
            'hip/hiprtc.h',
            'hipblas.h',
            'hiprand/hiprand.h',
            'hipsparse.h',
            'hipfft.h',
            'roctx.h',
            'rocsolver.h',
        ],
        'libraries': [
            'amdhip64',  # was hiprtc and hip_hcc before ROCm 3.8.0
            'hipblas',
            ('hipfft', lambda hip_version: hip_version >= 401),
            'hiprand',
            'hipsparse',
            'rocfft',
            'roctx64',
            'rocblas',
            'rocsolver',
            'rocsparse',
        ],
        'check_method': build.check_hip_version,
        'version_method': build.get_hip_version,
    }
    CUDA_cusolver = {
        'name': 'cusolver',
        'required': True,
        'file': [
            'cupy_backends.cuda.libs.cusolver',
            'cupyx.cusolver',
        ],
        'include': [
            'cusolverDn.h',
        ],
        'libraries': [
            'cusolver',
        ],
    }
    CUDA_cudnn = {
        'name': 'cudnn',
        'file': [
            'cupy_backends.cuda.libs.cudnn',
            'cupyx.cudnn',
        ],
        'include': [
            'cudnn.h',
        ],
        'libraries': [
            'cudnn',
        ],
        'check_method': build.check_cudnn_version,
        'version_method': build.get_cudnn_version,
    }
    CUDA_nccl = {
        'name': 'nccl',
        'file': [
            'cupy_backends.cuda.libs.nccl',
        ],
        'include': [
            'nccl.h',
        ],
        'libraries': [
            'nccl',
        ],
        'check_method': build.check_nccl_version,
        'version_method': build.get_nccl_version,
    }
    CUDA_nvtx = {
        'name': 'nvtx',
        'file': [
            'cupy_backends.cuda.libs.nvtx',
        ],
        'include': [
            'nvtx3/nvToolsExt.h',
        ],
        'libraries': [
        ],
        'check_method': build.check_nvtx,
    }
    CUDA_cutensor = {
        'name': 'cutensor',
        'file': [
            'cupy_backends.cuda.libs.cutensor',
            'cupyx.cutensor',
        ],
        'include': [
            'cutensor.h',
        ],
        'libraries': [
            'cutensor',
            'cublas',
        ],
        'check_method': build.check_cutensor_version,
        'version_method': build.get_cutensor_version,
    }
    CUDA_cub = {
        'name': 'cub',
        'required': True,
        'file': [
            ('cupy.cuda.cub', ['cupy/cuda/cupy_cub.cu']),
        ],
        'include': [
            'cub/util_namespace.cuh',  # dummy
        ],
        'libraries': list(_cudart_static_libs),
        'static_libraries': ['cudart_static'],
        'check_method': build.check_cub_version,
        'version_method': build.get_cub_version,
    }
    CUDA_jitify = {
        'name': 'jitify',
        'required': True,
        'file': [
            'cupy.cuda.jitify',
        ],
        'include': [
            'cuda.h',
            'cuda_runtime.h',
            'nvrtc.h',
        ],
        'libraries': [
            # Dependency from Jitify header files
            'cuda',
            'nvrtc',
        ] + _cudart_static_libs,
        'static_libraries': ['cudart_static'],
        'check_method': build.check_jitify_version,
        'version_method': build.get_jitify_version,
    }
    CUDA_random = {
        'name': 'random',
        'required': True,
        'file': [
            'cupy.random._bit_generator',
            ('cupy.random._generator_api',
             ['cupy/random/cupy_distributions.cu']),
        ],
        'include': [
        ],
        'libraries': [
            'curand',
        ] + _cudart_static_libs,
        'static_libraries': ['cudart_static'],
    }
    HIP_random = {
        'name': 'random',
        'required': True,
        'file': [
            'cupy.random._bit_generator',
            ('cupy.random._generator_api',
             ['cupy/random/cupy_distributions.cu']),
        ],
        'include': [
            'hiprand/hiprand.h',
        ],
        'libraries': [
            # Dependency from cuRAND header files
            'amdhip64',  # was hiprtc and hip_hcc before ROCm 3.8.0
            'hiprand',
        ],
        'check_method': build.check_hip_version,
        'version_method': build.get_hip_version,
    }
    CUDA_cusparselt = {
        'name': 'cusparselt',
        'file': [
            'cupy_backends.cuda.libs.cusparselt',
        ],
        'include': [
            'cusparseLt.h',
        ],
        'libraries': [
            'cusparseLt',
        ],
        'check_method': build.check_cusparselt_version,
        'version_method': build.get_cusparselt_version,
    }
    HIP_cub = {
        'name': 'cub',
        'required': True,
        'file': [
            ('cupy.cuda.cub', ['cupy/cuda/cupy_cub.cu']),
        ],
        'include': [
            'hipcub/hipcub_version.hpp',  # dummy
        ],
        'libraries': [
            'amdhip64',  # was hiprtc and hip_hcc before ROCm 3.8.0
        ],
        'check_method': build.check_cub_version,
        'version_method': build.get_cub_version,
    }
    HIP_nccl = {
        'name': 'nccl',
        'file': [
            'cupy_backends.cuda.libs.nccl',
        ],
        'include': [
            'rccl.h',
        ],
        'libraries': [
            'rccl',
        ],
        'check_method': build.check_nccl_version,
        'version_method': build.get_nccl_version,
    }
    HIP_thrust = {
        'name': 'thrust',
        'required': True,
        'file': [
            ('cupy.cuda.thrust', ['cupy/cuda/cupy_thrust.cu']),
        ],
        'include': [
            # WAR #9098:
            # rocThrust 3.3.0 (ROCm 6.4.0) cannot be compiled by host compiler
            # 'thrust/version.h',
        ],
        'libraries': [
            'amdhip64',  # was hiprtc and hip_hcc before ROCm 3.8.0
        ],
    }
    CUDA_thrust = {
        'name': 'thrust',
        'required': True,
        'file': [
            ('cupy.cuda.thrust', ['cupy/cuda/cupy_thrust.cu']),
        ],
        'include': [
            'thrust/version.h',
        ],
        'libraries': list(_cudart_static_libs),
        'static_libraries': ['cudart_static'],
        'check_method': build.check_thrust_version,
        'version_method': build.get_thrust_version,
    }
    COMMON_dlpack = {
        'name': 'dlpack',
        'required': True,
        'file': [
            'cupy._core.dlpack',
        ],
        'include': [
            'cupy/_dlpack/dlpack.h',
        ],
        'libraries': [],
    }

    if ctx.use_hip:
        features = [
            _from_dict(HIP_cuda_nvtx_cusolver, ctx),
            _from_dict(HIP_cub, ctx),
            _from_dict(HIP_nccl, ctx),
            _from_dict(HIP_random, ctx),
            _from_dict(HIP_thrust, ctx),
            _from_dict(COMMON_dlpack, ctx),
        ]
    else:
        features = [
            CUDA_cuda(ctx),
            _from_dict(CUDA_cusolver, ctx),
            _from_dict(CUDA_cudnn, ctx),
            _from_dict(CUDA_nccl, ctx),
            _from_dict(CUDA_nvtx, ctx),
            _from_dict(CUDA_cutensor, ctx),
            _from_dict(CUDA_cub, ctx),
            _from_dict(CUDA_jitify, ctx),
            _from_dict(CUDA_random, ctx),
            _from_dict(CUDA_thrust, ctx),
            _from_dict(CUDA_cusparselt, ctx),
            _from_dict(COMMON_dlpack, ctx),
        ]
    return {f.name: f for f in features}


class CUDA_cuda(Feature):
    minimum_cuda_version = 11020

    def __init__(self, ctx: Context):
        super().__init__(ctx)
        self.name = 'cuda'
        self.required = True
        self.modules = _cuda_files
        self.includes = [
            'cublas_v2.h',
            'cuda.h',
            'cuda_profiler_api.h',
            'cuda_runtime.h',
            'cufft.h',
            'curand.h',
            'cusparse.h',
        ]
        self.libraries = (
            # CUDA Runtime
            _cudart_static_libs +

            # CUDA Toolkit
            ['cublas', 'cufft', 'curand', 'cusparse']
        )
        self.static_libraries = ['cudart_static']
        self._version = self._UNDETERMINED

    def configure(self, compiler: Any, settings: Any) -> bool:
        try:
            out = build.build_and_run(compiler, '''
            #include <cuda.h>
            #include <stdio.h>
            int main() {
              printf("%d", CUDA_VERSION);
              return 0;
            }
            ''', include_dirs=settings['include_dirs'])  # type: ignore[no-untyped-call] # NOQA
        except Exception as e:
            utils.print_warning('Cannot check CUDA version', str(e))
            return False

        self._version = int(out)

        if self._version < self.minimum_cuda_version:
            utils.print_warning(
                'CUDA version is too old: %d' % self._version,
                'CUDA 11.2 or newer is required')
            return False
        return True
