from __future__ import annotations

from typing import Any, Callable, Final, List, Sequence, Type, Union
from triton.backends.nvidia.driver import CudaDriver
from triton.runtime import _allocation
from . import TVMFFILauncherImpl, utils, string_to_type


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

        self.signature: List[str] = [*src.signature.values()]
        self.num_ctas: Final[int] = getattr(metadata, "num_ctas", 1)
        self.global_scratch_size: Final[int] = metadata.global_scratch_size
        self.global_scratch_align: Final[int] = metadata.global_scratch_align
        self.profile_scratch_size: Final[int] = metadata.profile_scratch_size
        self.profile_scratch_align: Final[int] = metadata.profile_scratch_align
        self.launch_cooperative_grid: Final[bool] = metadata.launch_cooperative_grid
        self.launch_pdl: Final[bool] = metadata.launch_pdl
        self.impl: TVMFFILauncherImpl = TVMFFILauncherImpl(
            [string_to_type(t) for t in self.signature],
            self.launch_cooperative_grid,
            self.launch_pdl,
        )
        self.launch: Callable[
            [
                int,
                int,
                int,
                int,
                int,
                tuple[int, int, int],
                object,
                object,
                object,
                object,
                object,
                Sequence[Union[Any]],
            ]
        ] = self.impl.launch

    def __call__(
        self,
        gridX,
        gridY,
        gridZ,
        stream,
        function,
        kernel_metadata,
        launch_metadata,
        launch_enter_hook,
        launch_exit_hook,
        *args,
    ):
        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

        # args: Sequence[TypedValue] = TypedValue.make_typed_values(self.signature, args)

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


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


del CudaDriver