from __future__ import annotations

from ctypes import c_void_p
from typing import Any, List, Mapping, Optional, Tuple, Type
from triton.backends.nvidia.driver import CudaDriver
from . import TypedValue, get_device_properties, launch, load_binary, string_to_type


class TVMFFIUtils(object):
    def __new__(cls: Type[TVMFFIUtils]) -> TVMFFIUtils:
        if not hasattr(cls, "instance"):
            cls.instance = super().__new__(cls)
        return cls.instance

    def __init__(self, *args, **kwargs) -> None:
        super().__init__(*args, **kwargs)

    def load_binary(
        self, name: str, data: bytes, shared: int, device: int
    ) -> Tuple[c_void_p, c_void_p, int, int, int]:
        return load_binary(name, data, shared, device)

    def get_device_properties(self, device_id: int) -> Mapping[str, int]:
        return get_device_properties(device_id)

    def cuOccupancyMaxActiveClusters(self, *args, **kwargs):
        raise NotImplementedError(
            '"cuOccupancyMaxActiveClusters isn\'t expected to be invoked"'
        )

    def set_printf_fifo_size(self, *args, **kwargs):
        raise NotImplementedError(
            '"set_printf_fifo_size" isn\'t expected to be invoked'
        )

    def fill_tma_descriptor(self, *args, **kwargs):
        raise NotImplementedError(
            '"fill_tma_descriptor" hasn\'t been supported for Hopper'
        )

    def launch(self, *args, **kwargs):
        raise NotImplementedError(
            '"launch" is introduced in triton after commit d2b3925410689155e0f6028e8554bba972989348, which is still not supported yed'
        )

    def build_signature_metadata(self, *args, **kwargs):
        raise NotImplementedError(
            '"launch" is introduced in triton after commit d2b3925410689155e0f6028e8554bba972989348, which is still not supported yed'
        )


class TVMLauncher(object):
    def __init__(self, src: List[bool], metadata, *args, **kwargs) -> TVMLauncher:
        super().__init__(*args, **kwargs)

        self.signature = src.signature.values()
        self.num_ctas = getattr(metadata, "num_ctas", 1)
        self.launch = launch
        self.global_scratch_size = metadata.global_scratch_size
        self.global_scratch_align = metadata.global_scratch_align
        self.profile_scratch_size = metadata.profile_scratch_size
        self.profile_scratch_align = metadata.profile_scratch_align
        self.launch_cooperative_grid = metadata.launch_cooperative_grid
        self.launch_pdl = metadata.launch_pdl

    def __call__(
        self,
        gridX,
        gridY,
        gridZ,
        stream,
        function,
        kernel_metadata,
        launch_metadata,
        launch_enter_hook,
        launch_exit_hook,
        *args,
    ):
        from triton.runtime import _allocation

        def allocate_scratch(size, align, allocator):
            if size > 0:
                grid_size = gridX * gridY * gridZ
                alloc_size = grid_size * self.num_ctas * size
                alloc_fn = allocator.get()
                return alloc_fn(alloc_size, align, stream)
            return None

        global_scratch = allocate_scratch(
            self.global_scratch_size, self.global_scratch_align, _allocation._allocator
        )
        profile_scratch = allocate_scratch(
            self.profile_scratch_size,
            self.profile_scratch_align,
            _allocation._profile_allocator,
        )
        assert not self.launch_cooperative_grid
        assert not self.launch_pdl
        assert len(self.signature) == len(args)

        def canonicalize(arg: Any, sig: str) -> TypedValue:
            ty: Optional[int] = string_to_type(sig)
            assert ty is not None, sig
            return TypedValue(ty, arg)

        args = [canonicalize(arg, sig) for arg, sig in zip(args, self.signature)]

        return launch(
            gridX,
            gridY,
            gridZ,
            stream,
            function,
            kernel_metadata,
            launch_metadata,
            launch_enter_hook,
            launch_exit_hook,
            self.launch_cooperative_grid,
            self.launch_pdl,
            global_scratch,
            profile_scratch,
            *args,
        )


class TVMFFIDriver(CudaDriver):
    def __init__(self, *args, **kwargs) -> TVMFFIDriver:
        super().__init__(*args, **kwargs)
        self.utils: TVMFFIUtils = TVMFFIUtils()
        self.launcher_cls: Type[TVMLauncher] = TVMLauncher


del CudaDriver