from __future__ import annotations

from ctypes import c_void_p
from typing import Mapping, Tuple, Type
from triton.backends.nvidia.driver import CudaDriver
from .utils import get_device_properties, load_binary


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 TVMFFIDriver(CudaDriver):
    def __init__(self, *args, **kwargs) -> TVMFFIDriver:
        super().__init__(*args, **kwargs)
        self.utils: TVMFFIUtils = TVMFFIUtils()


del CudaDriver