import ctypes

_c_str_dltensor = b"dltensor"


class DLDeviceType(ctypes.c_int):
    """The enum that encodes the type of the device where
    DLTensor memory is allocated.
    """

    kDLCPU = 1
    kDLCUDA = 2
    kDLCUDAHost = 3
    kDLOpenCL = 4
    kDLVulkan = 7
    kDLMetal = 8
    kDLVPI = 9
    kDLROCM = 10
    kDLROCMHost = 11
    kDLCUDAManaged = 13
    kDLOneAPI = 14

    def __str__(self):
        return {
            self.kDLCPU: "CPU",
            self.kDLCUDA: "CUDA",
            self.kDLCUDAHost: "CUDAHost",
            self.kDLOpenCL: "OpenCL",
            self.kDLVulkan: "Vulkan",
            self.kDLMetal: "Metal",
            self.kDLVPI: "VPI",
            self.kDLROCM: "ROCM",
            self.kDLROCMHost: "ROMCHost",
            self.kDLCUDAManaged: "CUDAManaged",
            self.kDLOneAPI: "oneAPI",
        }[self.value]


class DLDevice(ctypes.Structure):
    """Represents the device where DLTensor memory is allocated.
    The device is represented by the pair of fields:
       device_type: DLDeviceType
       device_id: c_int
    """

    _fields_ = [
        ("device_type", DLDeviceType),
        ("device_id", ctypes.c_int),
    ]


class DLDataTypeCode(ctypes.c_uint8):
    """An integer that encodes the category of DLTensor elements' data type."""

    kDLInt = 0
    kDLUInt = 1
    kDLFloat = 2
    kDLOpaquePointer = 3
    kDLBfloat = 4
    kDLComplex = 5

    def __str__(self):
        return {
            self.kDLInt: "int",
            self.kDLUInt: "uint",
            self.kDLFloat: "float",
            self.kDLBfloat: "bfloat",
            self.kDLComplex: "complex",
            self.kDLOpaquePointer: "void_p",
        }[self.value]


class DLDataType(ctypes.Structure):
    """Descriptor of data type for elements of DLTensor.
    The data type is described by a triple, `DLDataType.type_code`,
    `DLDataType.bits`, and `DLDataType.lanes`.

    The element is understood as packed `lanes` repetitions of
    elements from `type_code` data-category of width `bits`.
    """

    _fields_ = [
        ("type_code", DLDataTypeCode),
        ("bits", ctypes.c_uint8),
        ("lanes", ctypes.c_uint16),
    ]
    TYPE_MAP = {
        "bool": (DLDataTypeCode.kDLUInt, 1, 1),
        "int8": (DLDataTypeCode.kDLInt, 8, 1),
        "int16": (DLDataTypeCode.kDLInt, 16, 1),
        "int32": (DLDataTypeCode.kDLInt, 32, 1),
        "int64": (DLDataTypeCode.kDLInt, 64, 1),
        "uint8": (DLDataTypeCode.kDLUInt, 8, 1),
        "uint16": (DLDataTypeCode.kDLUInt, 16, 1),
        "uint32": (DLDataTypeCode.kDLUInt, 32, 1),
        "uint64": (DLDataTypeCode.kDLUInt, 64, 1),
        "float16": (DLDataTypeCode.kDLFloat, 16, 1),
        "float32": (DLDataTypeCode.kDLFloat, 32, 1),
        "float64": (DLDataTypeCode.kDLFloat, 64, 1),
        "complex64": (DLDataTypeCode.kDLComplex, 64, 1),
        "complex128": (DLDataTypeCode.kDLComplex, 128, 1),
    }


class DLTensor(ctypes.Structure):
    """Structure describing strided layout of DLTensor.
    Fields are:
       data:  void pointer
       device: DLDevice
       ndim: number of indices needed to reference an
             element of the tensor
       dtype: data type descriptor
       shape: tuple with lengths of the corresponding
              tensor dimensions
       strides: tuple of numbers of array elements to
                step in each dimension when traversing
                the tensor
       byte_offset: data + byte_offset gives the address of
                tensor element with index (0,) * ndim
    """

    _fields_ = [
        ("data", ctypes.c_void_p),
        ("device", DLDevice),
        ("ndim", ctypes.c_int),
        ("dtype", DLDataType),
        ("shape", ctypes.POINTER(ctypes.c_int64)),
        ("strides", ctypes.POINTER(ctypes.c_int64)),
        ("byte_offset", ctypes.c_uint64),
    ]


class DLManagedTensor(ctypes.Structure):
    """Structure storing the pointer to the tensor descriptor,
    deleter callable for the tensor descriptor, and pointer to
    some additional data. These are stored in fields `dl_tensor`,
    `deleter`, and `manager_ctx`."""

    _fields_ = [
        ("dl_tensor", DLTensor),
        ("manager_ctx", ctypes.c_void_p),
        ("deleter", ctypes.CFUNCTYPE(None, ctypes.c_void_p)),
    ]