from __future__ import annotations

from typing import Any, List, Optional, Type
from triton.backends.nvidia.driver import CudaDriver
from triton.runtime import _allocation
from . import TypedValue, 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: int = getattr(metadata, "num_ctas", 1)
        self.launch = utils.launch
        self.global_scratch_size: int = metadata.global_scratch_size
        self.global_scratch_align: int = metadata.global_scratch_align
        self.profile_scratch_size: int = metadata.profile_scratch_size
        self.profile_scratch_align: int = metadata.profile_scratch_align
        self.launch_cooperative_grid: bool = metadata.launch_cooperative_grid
        self.launch_pdl: bool = metadata.launch_pdl

    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
        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 self.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 = utils
        self.launcher_cls: Type[TVMLauncher] = TVMLauncher


del CudaDriver