from __future__ import annotations

from typing import Mapping, Type
from triton.backends.nvidia.driver import CudaDriver
from .utils import get_device_properties


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)
        from triton.backends.nvidia.driver import CudaUtils

        self._utils: CudaUtils = CudaUtils()

    def load_binary(self, *args, **kwargs):
        return self._utils.load_binary(*args, **kwargs)

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

    def cuOccupancyMaxActiveClusters(self, *args, **kwargs):
        return self._utils.cuOccupancyMaxActiveClusters(*args, **kwargs)

    def set_printf_fifo_size(self, *args, **kwargs):
        return self._utils.set_printf_fifo_size(*args, **kwargs)

    def fill_tma_descriptor(self, *args, **kwargs):
        return self._utils.fill_tma_descriptor(*args, **kwargs)

    def launch(self, *args, **kwargs):
        return self._utils.launch(*args, **kwargs)

    def build_signature_metadata(self, *args, **kwargs):
        return self._utils.build_signature_metadata(*args, **kwargs)


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


del CudaDriver