cimport cpython  # NOQA

from libc cimport stdlib
from libc.stdint cimport uint8_t
from libc.stdint cimport uint16_t
from libc.stdint cimport int64_t
from libc.stdint cimport uint64_t
from libc.stdint cimport intptr_t
from libcpp.vector cimport vector

from cupy_backends.cuda.api cimport runtime
from cupy_backends.cuda cimport stream as stream_module
from cupy._core.core cimport ndarray
from cupy.cuda cimport device
from cupy.cuda cimport memory

import warnings

import cupy


cdef extern from './include/cupy/dlpack/dlpack.h' nogil:
    '''
    // stringification is currently needed as the version is "030"...
    #define _stringify_(s) #s
    #define _xstringify_(s) _stringify_(s)
    const char* DLPACK_VER_STR = _xstringify_(DLPACK_VERSION);
    #undef _xstringify_
    #undef _stringify_
    '''
    const char* DLPACK_VER_STR

    cdef enum DLDeviceType:
        kDLCPU
        kDLGPU
        kDLCPUPinned
        kDLOpenCL
        kDLVulkan
        kDLMetal
        kDLVPI
        kDLROCM

    ctypedef struct DLDevice:
        DLDeviceType device_type
        int device_id

    cdef enum DLDataTypeCode:
        kDLInt
        kDLUInt
        kDLFloat
        kDLBfloat
        kDLComplex

    ctypedef struct DLDataType:
        uint8_t code
        uint8_t bits
        uint16_t lanes

    ctypedef struct DLTensor:
        void* data
        DLDevice device
        int ndim
        DLDataType dtype
        int64_t* shape
        int64_t* strides
        uint64_t byte_offset

    ctypedef struct DLManagedTensor:
        DLTensor dl_tensor
        void* manager_ctx
        void (*deleter)(DLManagedTensor*)  # noqa: E211


def get_build_version():
    return DLPACK_VER_STR.decode()


cdef void pycapsule_deleter(object dltensor):
    cdef DLManagedTensor* dlm_tensor
    # Do not invoke the deleter on a used capsule
    if cpython.PyCapsule_IsValid(dltensor, 'dltensor'):
        dlm_tensor = <DLManagedTensor*>cpython.PyCapsule_GetPointer(
            dltensor, 'dltensor')
        dlm_tensor.deleter(dlm_tensor)


cdef void deleter(DLManagedTensor* tensor) with gil:
    if tensor.manager_ctx is NULL:
        return
    stdlib.free(tensor.dl_tensor.shape)
    cpython.Py_DECREF(<ndarray>tensor.manager_ctx)
    tensor.manager_ctx = NULL
    stdlib.free(tensor)


# The name of this function is following the framework integration guide of
# TensorComprehensions.
cpdef object toDlpack(ndarray array) except +:
    cdef DLManagedTensor* dlm_tensor = \
        <DLManagedTensor*>stdlib.malloc(sizeof(DLManagedTensor))

    cdef size_t ndim = array._shape.size()
    cdef DLTensor* dl_tensor = &dlm_tensor.dl_tensor
    dl_tensor.data = <void*>array.data.ptr
    dl_tensor.ndim = ndim

    cdef int64_t* shape_strides = \
        <int64_t*>stdlib.malloc(ndim * sizeof(int64_t) * 2)
    for n in range(ndim):
        shape_strides[n] = array._shape[n]
    dl_tensor.shape = shape_strides
    for n in range(ndim):
        shape_strides[n + ndim] = array._strides[n] // array.dtype.itemsize

    dl_tensor.strides = shape_strides + ndim
    dl_tensor.byte_offset = 0

    cdef DLDevice* device = &dl_tensor.device
    # TODO(leofang): if the CuPy array is backed by managed memory, we should
    # switch to kDLCUDAManaged
    if not runtime._is_hip_environment:
        device.device_type = kDLGPU
    else:
        device.device_type = kDLROCM
    device.device_id = array.data.device_id

    cdef DLDataType* dtype = &dl_tensor.dtype
    if array.dtype.kind == 'u':
        dtype.code = <uint8_t>kDLUInt
    elif array.dtype.kind == 'i':
        dtype.code = <uint8_t>kDLInt
    elif array.dtype.kind == 'f':
        dtype.code = <uint8_t>kDLFloat
    elif array.dtype.kind == 'c':
        dtype.code = <uint8_t>kDLComplex
    else:
        raise ValueError('Unknown dtype')
    dtype.lanes = <uint16_t>1
    dtype.bits = <uint8_t>(array.dtype.itemsize * 8)

    dlm_tensor.manager_ctx = <void*>array
    cpython.Py_INCREF(array)
    dlm_tensor.deleter = deleter

    return cpython.PyCapsule_New(dlm_tensor, 'dltensor', pycapsule_deleter)


# TODO(leofang): Support kDLCUDAPinned, kDLCUDAManaged, and kDLROCMPinned
cdef class DLPackMemory(memory.BaseMemory):

    """Memory object for a dlpack tensor.

    This does not allocate any memory.

    """

    cdef DLManagedTensor* dlm_tensor
    cdef object dltensor

    def __init__(self, object dltensor):
        cdef DLManagedTensor* dlm_tensor

        # sanity checks
        if not cpython.PyCapsule_IsValid(dltensor, 'dltensor'):
            raise ValueError('A DLPack tensor object cannot be consumed '
                             'multiple times')
        dlm_tensor = <DLManagedTensor*>cpython.PyCapsule_GetPointer(
            dltensor, 'dltensor')
        if runtime._is_hip_environment:
            if dlm_tensor.dl_tensor.device.device_type != kDLROCM:
                raise RuntimeError('CuPy is built against ROCm/HIP, different '
                                   'from the backend that backs the incoming '
                                   'DLPack tensor')
        else:
            if dlm_tensor.dl_tensor.device.device_type != kDLGPU:
                raise RuntimeError('CuPy is built against CUDA, different '
                                   'from the backend that backs the incoming '
                                   'DLPack tensor')

        self.dltensor = dltensor
        self.dlm_tensor = dlm_tensor
        self.device_id = dlm_tensor.dl_tensor.device.device_id
        self.ptr = <intptr_t>dlm_tensor.dl_tensor.data
        cdef int n = 0, s = 0
        cdef int ndim = dlm_tensor.dl_tensor.ndim
        cdef int64_t* shape = dlm_tensor.dl_tensor.shape
        for s in shape[:ndim]:
            n += s
        self.size = dlm_tensor.dl_tensor.dtype.bits * n // 8

    def __dealloc__(self):
        cdef DLManagedTensor* dlm_tensor = self.dlm_tensor
        # dlm_tensor could be uninitialized if an error is raised in __init__
        if dlm_tensor != NULL:
            dlm_tensor.deleter(dlm_tensor)


# The name of this function is following the framework integration guide of
# TensorComprehensions.
cpdef ndarray fromDlpack(object dltensor) except +:
    """Zero-copy conversion from a DLPack tensor to a :class:`~cupy.ndarray`.

    DLPack is a open in memory tensor structure proposed in this repository:
    `dmlc/dlpack <https://github.com/dmlc/dlpack>`_.

    This function takes a :class:`PyCapsule` object which contains a pointer to
    a DLPack tensor as input, and returns a :class:`~cupy.ndarray`. This
    function does not copy the data in the DLPack tensor but both
    DLPack tensor and :class:`~cupy.ndarray` have pointers which are pointing
    to the same memory region for the data.

    Args:
        dltensor (:class:`PyCapsule`): Input DLPack tensor which is
            encapsulated in a :class:`PyCapsule` object.

    Returns:
        array (:class:`~cupy.ndarray`): A CuPy ndarray.

    .. warning::

        This function is deprecated in favor of :func:`~cupy.from_dlpack` and
        will be removed in a future version of CuPy.

    .. warning::

        As of the DLPack v0.5 specification, it is implicitly assumed that
        the user is responsible to ensure the Producer and the Consumer are
        operating on the same stream.

    .. seealso::

        :meth:`cupy.ndarray.toDlpack` is a method for zero-copy conversion
        from a :class:`~cupy.ndarray` to a DLPack tensor (which is encapsulated
        in a :class:`PyCapsule` object).

    .. admonition:: Example

        >>> import cupy
        >>> array1 = cupy.array([0, 1, 2], dtype=cupy.float32)
        >>> dltensor = array1.toDlpack()
        >>> array2 = cupy.fromDlpack(dltensor)
        >>> cupy.testing.assert_array_equal(array1, array2)

    """
    warnings.warn('This function is deprecated in favor of cupy.from_dlpack',
                  DeprecationWarning)
    return _dlpack_to_cupy_array(dltensor)


cdef inline ndarray _dlpack_to_cupy_array(dltensor) except +:
    cdef DLPackMemory mem = DLPackMemory(dltensor)
    cdef DLDataType dtype = mem.dlm_tensor.dl_tensor.dtype
    cdef int bits = dtype.bits
    if dtype.lanes != 1:
        raise ValueError(f'vector dtypes (lanes={dtype.lanes}) is '
                         'not supported')
    if dtype.code == kDLUInt:
        if bits == 8:
            cp_dtype = cupy.uint8
        elif bits == 16:
            cp_dtype = cupy.uint16
        elif bits == 32:
            cp_dtype = cupy.uint32
        elif bits == 64:
            cp_dtype = cupy.uint64
        else:
            raise TypeError('uint{} is not supported.'.format(bits))
    elif dtype.code == kDLInt:
        if bits == 8:
            cp_dtype = cupy.int8
        elif bits == 16:
            cp_dtype = cupy.int16
        elif bits == 32:
            cp_dtype = cupy.int32
        elif bits == 64:
            cp_dtype = cupy.int64
        else:
            raise TypeError('int{} is not supported.'.format(bits))
    elif dtype.code == kDLFloat:
        if bits == 16:
            cp_dtype = cupy.float16
        elif bits == 32:
            cp_dtype = cupy.float32
        elif bits == 64:
            cp_dtype = cupy.float64
        else:
            raise TypeError('float{} is not supported.'.format(bits))
    elif dtype.code == kDLComplex:
        # TODO(leofang): support complex32
        if bits == 64:
            cp_dtype = cupy.complex64
        elif bits == 128:
            cp_dtype = cupy.complex128
        else:
            raise TypeError('complex{} is not supported.'.format(bits))
    elif dtype.code == kDLBfloat:
        raise NotImplementedError('CuPy does not support bfloat16 yet')
    else:
        raise TypeError('Unsupported dtype. dtype code: {}'.format(dtype.code))

    mem_ptr = memory.MemoryPointer(mem, mem.dlm_tensor.dl_tensor.byte_offset)
    cdef int64_t ndim = mem.dlm_tensor.dl_tensor.ndim

    cdef int64_t* shape = mem.dlm_tensor.dl_tensor.shape
    cdef vector[Py_ssize_t] shape_vec
    shape_vec.assign(shape, shape + ndim)

    if mem.dlm_tensor.dl_tensor.strides is NULL:
        # Make sure this capsule will never be used again.
        cpython.PyCapsule_SetName(mem.dltensor, 'used_dltensor')
        return ndarray(shape_vec, cp_dtype, mem_ptr, strides=None)
    cdef int64_t* strides = mem.dlm_tensor.dl_tensor.strides
    cdef vector[Py_ssize_t] strides_vec
    for i in range(ndim):
        strides_vec.push_back(strides[i] * (bits // 8))

    # Make sure this capsule will never be used again.
    cpython.PyCapsule_SetName(mem.dltensor, 'used_dltensor')
    return ndarray(shape_vec, cp_dtype, mem_ptr, strides=strides_vec)


# TODO(leofang): this function is exposed to the cupy namespace, so it returns
# a cupy.ndarray which is not compliant with the Python array API. When we have
# a compliant object living in, say, cupy.array_api, we will expose another
# function cupy.array_api.from_dlpack().
cpdef from_dlpack(array):
    """Zero-copy conversion between array objects compliant with the DLPack
    data exchange protocol.

    Args:
        array (object): an array object that implements two methods:
            ``__dlpack__()`` and ``__dlpack_device__()``.

    Returns:
        cupy.ndarray: a CuPy array that can be safely accessed on CuPy's
        current stream.

    .. note::
        This function is different from CuPy's legacy :func:`~cupy.fromDlpack`
        function. This function takes any object implementing the DLPack data
        exchange protocol, as well as a raw :class:`PyCapsule` object that
        contains the DLPack tensor as input (for backward compatibility),
        whereas :func:`~cupy.fromDlpack` only accepts :class:`PyCapsule`
        objects. If the input object is not compliant with the protocol, users
        are responsible to ensure data safety.

    .. seealso::
        `Data interchange mechanisms`_

    .. _Data interchange mechanisms:
        https://data-apis.org/array-api/latest/design_topics/data_interchange.html
    """
    if not hasattr(array, '__dlpack_device__'):
        # backward compatibility: accept passing in a pycapsule
        dltensor = array
        return _dlpack_to_cupy_array(dltensor)
    else:
        dev_type, dev_id = array.__dlpack_device__()

    # CuPy is the consumer, so we provide our current stream to the producer
    if dev_type == <int>kDLGPU:
        with device.Device(dev_id):
            assert not runtime._is_hip_environment
            stream = stream_module.get_current_stream_ptr()
            if stream == 0:
                stream = stream_module.get_default_stream_ptr()
            dltensor = array.__dlpack__(stream=stream)
    elif dev_type == <int>kDLROCM:
        with device.Device(dev_id):
            assert runtime._is_hip_environment
            stream = stream_module.get_current_stream_ptr()
            dltensor = array.__dlpack__(stream=stream)
    else:
        # TODO(leofang): support kDLCUDAPinned, kDLCUDAManaged, etc
        dltensor = None
        raise ValueError

    return _dlpack_to_cupy_array(dltensor)
