# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.

# cython: language_level = 3

from libc.stdint cimport uintptr_t, int8_t, uint8_t, uint16_t, int64_t
from libc.string cimport memcpy
from libc.stdio cimport snprintf
from cpython.pycapsule cimport PyCapsule_New, PyCapsule_GetPointer, PyCapsule_IsValid
from cpython cimport (
    Py_buffer,
    PyObject_GetBuffer,
    PyBuffer_Release,
    PyBuffer_ToContiguous,
    PyBuffer_FillInfo,
    PyBUF_ANY_CONTIGUOUS,
    PyBUF_FORMAT,
    PyBUF_WRITABLE,
    PyErr_WriteUnraisable,
)
from cpython.ref cimport Py_INCREF, Py_DECREF

from nanoarrow_c cimport (
    NANOARROW_OK,
    ArrowMalloc,
    ArrowFree,
    ArrowType,
    ArrowTypeString,
    ArrowBitGet,
    ArrowBitsUnpackInt8,
    ArrowBufferReserve,
    ArrowBufferAppendFill,
    ArrowBufferAppendInt8,
    ArrowBitmapInit,
    ArrowBitmapReset,
    ArrowBitmap,
    ArrowBitmapReserve,
    ArrowBitmapAppend,
    ArrowBitmapAppendUnsafe,
    ArrowBuffer,
    ArrowBufferMove,
)

from nanoarrow_device_c cimport (
    ARROW_DEVICE_CPU,
    ARROW_DEVICE_CUDA,
    ArrowDevice,
)

from nanoarrow_dlpack cimport (
    DLDataType,
    DLDevice,
    DLDeviceType,
    DLManagedTensor,
    DLTensor,
    kDLCPU,
    kDLFloat,
    kDLInt,
    kDLUInt
)

from nanoarrow cimport _utils
from nanoarrow cimport _types
from nanoarrow._device cimport CSharedSyncEvent, Device

from struct import unpack_from, iter_unpack, calcsize, Struct

from nanoarrow import _repr_utils
from nanoarrow._device import DEVICE_CPU


cdef void pycapsule_dlpack_deleter(object dltensor) noexcept:
    cdef DLManagedTensor* dlm_tensor

    # Do nothing if the capsule has been consumed
    if PyCapsule_IsValid(dltensor, "used_dltensor"):
        return

    dlm_tensor = <DLManagedTensor*>PyCapsule_GetPointer(dltensor, 'dltensor')
    if dlm_tensor == NULL:
        PyErr_WriteUnraisable(dltensor)
    # The deleter can be NULL if there is no way for the caller
    # to provide a reasonable destructor
    elif dlm_tensor.deleter:
        dlm_tensor.deleter(dlm_tensor)


cdef void view_dlpack_deleter(DLManagedTensor* tensor) noexcept with gil:
    if tensor.manager_ctx is NULL:
        return
    Py_DECREF(<CBufferView>tensor.manager_ctx)
    tensor.manager_ctx = NULL
    ArrowFree(tensor)


cdef DLDataType view_to_dlpack_data_type(CBufferView view):
    cdef DLDataType dtype
    # Define DLDataType struct
    if _types.is_unsigned_integer(view.data_type_id):
        dtype.code = kDLUInt
    elif _types.is_signed_integer(view.data_type_id):
        dtype.code = kDLInt
    elif _types.is_floating_point(view.data_type_id):
        dtype.code = kDLFloat
    elif _types.equal(view.data_type_id, _types.BOOL):
        raise ValueError('Bit-packed boolean data type not supported by DLPack.')
    else:
        raise ValueError('DataType is not compatible with DLPack spec: ' + view.data_type)
    dtype.lanes = <uint16_t>1
    dtype.bits = <uint8_t>(view._element_size_bits)

    return dtype

cdef int dlpack_data_type_to_arrow(DLDataType dtype):
    if dtype.code == kDLInt:
        if dtype.bits == 8:
            return _types.INT8
        elif dtype.bits == 16:
            return _types.INT16
        elif dtype.bits == 32:
            return _types.INT32
        elif dtype.bits == 64:
            return _types.INT64
    elif dtype.code == kDLUInt:
        if dtype.bits == 8:
            return _types.UINT8
        elif dtype.bits == 16:
            return _types.UINT16
        elif dtype.bits == 32:
            return _types.UINT32
        elif dtype.bits == 64:
            return _types.UINT64
    elif dtype.code == kDLFloat:
        if dtype.bits == 16:
            return _types.HALF_FLOAT
        elif dtype.bits == 32:
            return _types.FLOAT
        elif dtype.bits == 64:
            return _types.DOUBLE

    raise ValueError("Can't convert dlpack data type to Arrow type")

cdef object view_to_dlpack(CBufferView view, stream=None):
    # Define DLDevice and DLDataType struct and
    # with that check for data type support first
    cdef DLDevice device = view_to_dlpack_device(view)
    cdef DLDataType dtype = view_to_dlpack_data_type(view)

    # Allocate memory for DLManagedTensor
    cdef DLManagedTensor* dlm_tensor = <DLManagedTensor*>ArrowMalloc(sizeof(DLManagedTensor))
    # Define DLManagedTensor struct
    cdef DLTensor* dl_tensor = &dlm_tensor.dl_tensor
    dl_tensor.data = <void*>view._ptr.data.data
    dl_tensor.ndim = 1

    dl_tensor.shape = &view._n_elements
    dl_tensor.strides = NULL
    dl_tensor.byte_offset = 0

    dl_tensor.device = device
    dl_tensor.dtype = dtype

    dlm_tensor.manager_ctx = <void*>view
    Py_INCREF(view)
    dlm_tensor.deleter = view_dlpack_deleter

    # stream has a DLPack + device specific interpretation

    # nanoarrow_device needs a CUstream* (i.e., a CUstream_st**), but dlpack
    # gives us a CUstream_st*.
    cdef void* cuda_pstream

    if view._event.device is DEVICE_CPU:
        if stream is not None and stream != -1:
            raise ValueError("dlpack stream must be None or -1 for the CPU device")
    elif view._event.device.device_type_id == ARROW_DEVICE_CUDA:
        if stream == 0:
            raise ValueError("dlpack stream value of 0 is not permitted for CUDA")
        elif stream == -1:
            # Sentinel for "do not synchronize"
            pass
        elif stream in (1, 2):
            # Technically we are mixing the per-thread and legacy default streams here;
            # however, the nanoarrow_device API currently has no mechanism to expose
            # a pointer to these streams specifically.
            cuda_pstream = <void*>0
            view._event.synchronize_stream(<uintptr_t>&cuda_pstream)
        else:
            # Otherwise, this is a CUstream** (i.e., CUstream_st*)
            cuda_pstream = <void*><uintptr_t>stream
            view._event.synchronize_stream(<uintptr_t>&cuda_pstream)

    return PyCapsule_New(dlm_tensor, 'dltensor', pycapsule_dlpack_deleter)


cdef DLDevice view_to_dlpack_device(CBufferView view):
    cdef DLDevice device

    # Check data type support
    if _types.equal(view.data_type_id, _types.BOOL):
        raise ValueError('Bit-packed boolean data type not supported by DLPack.')
    elif (
        not _types.is_unsigned_integer(view.data_type_id)
        and not _types.is_signed_integer(view.data_type_id)
        and not _types.is_floating_point(view.data_type_id)
    ):
        raise ValueError('DataType is not compatible with DLPack spec: ' + view.data_type)

    # Define DLDevice struct
    cdef ArrowDevice* arrow_device = view._event.device._ptr
    if arrow_device.device_type is ARROW_DEVICE_CPU:
        # DLPack uses 0 for the CPU device id where Arrow uses -1
        device.device_type = kDLCPU
        device.device_id =  0
    else:
        # Otherwise, Arrow's device identifiers and types are intentionally
        # identical to DLPack
        device.device_type = <DLDeviceType>arrow_device.device_type
        device.device_id = arrow_device.device_id

    return device


cdef bint dlpack_strides_are_contiguous(DLTensor* dl_tensor):
    if dl_tensor.strides == NULL:
        return True

    if dl_tensor.ndim != 1:
        raise NotImplementedError("Contiguous stride check not implemented for ndim != 1")

    # DLTensor strides are in elemements, not bytes
    return dl_tensor.strides[0] == 1


cdef class CBufferView:
    """Wrapper for Array buffer content

    This object is a Python wrapper around a buffer held by an Array.
    It implements the Python buffer protocol and is best accessed through
    another implementor (e.g., `np.array(array_view.buffers[1])`)). Note that
    this buffer content does not apply any parent offset.
    """

    def __cinit__(self, object base, uintptr_t addr, int64_t size_bytes,
                  ArrowType data_type,
                  Py_ssize_t element_size_bits, CSharedSyncEvent event):
        self._base = base
        self._ptr.data.data = <void*>addr
        self._ptr.size_bytes = size_bytes
        self._data_type = data_type
        self._event = event
        self._format[0] = 0
        self._element_size_bits = _types.to_format(
            self._data_type,
            element_size_bits,
            sizeof(self._format),
            self._format
        )
        self._strides = self._item_size()
        self._shape = self._ptr.size_bytes // self._strides
        if _types.equal(self._data_type, _types.BOOL):
            self._n_elements = self._shape * 8
        else:
            self._n_elements = self._shape

    def _addr(self):
        return <uintptr_t>self._ptr.data.data

    @property
    def device(self):
        return self._event.device

    @property
    def element_size_bits(self):
        return self._element_size_bits

    @property
    def size_bytes(self):
        return self._ptr.size_bytes

    @property
    def data_type_id(self):
        return self._data_type

    @property
    def data_type(self):
        return ArrowTypeString(self._data_type).decode("UTF-8")

    @property
    def format(self):
        return self._format.decode("UTF-8")

    @property
    def itemsize(self):
        return self._strides

    def __len__(self):
        return self._shape

    def __getitem__(self, int64_t i):
        if i < 0 or i >= self._shape:
            raise IndexError(f"Index {i} out of range")
        cdef int64_t offset = self._strides * i
        value = unpack_from(self.format, buffer=self, offset=offset)
        if len(value) == 1:
            return value[0]
        else:
            return value

    def __iter__(self):
        return self._iter_dispatch(0, len(self))

    def _iter_dispatch(self, int64_t offset, int64_t length):
        if offset < 0 or length < 0 or (offset + length) > len(self):
            raise IndexError(
                f"offset {offset} and length {length} do not describe a valid slice "
                f"of buffer with length {len(self)}"
            )
        # memoryview's implementation is very fast but not always possible (half float, fixed-size binary, interval)
        if _types.one_of(
            self._data_type,
            (
                _types.HALF_FLOAT,
                _types.INTERVAL_DAY_TIME,
                _types.INTERVAL_MONTH_DAY_NANO,
                _types.DECIMAL128,
                _types.DECIMAL256
            )
        ) or (
            _types.equal(self._data_type, _types.BINARY) and self._element_size_bits != 0
        ):
            return self._iter_struct(offset, length)
        else:
            return self._iter_memoryview(offset, length)

    def _iter_memoryview(self, int64_t offset, int64_t length):
        return iter(memoryview(self)[offset:(offset + length)])

    def _iter_struct(self, int64_t offset, int64_t length):
        for value in iter_unpack(self.format, self):
            if len(value) == 1:
                yield value[0]
            else:
                yield value

    @property
    def n_elements(self):
        return self._n_elements

    def element(self, i):
        if _types.equal(self._data_type, _types.BOOL):
            if i < 0 or i >= self.n_elements:
                raise IndexError(f"Index {i} out of range")
            return ArrowBitGet(self._ptr.data.as_uint8, i)
        else:
            return self[i]

    def elements(self, offset=0, length=None):
        if length is None:
            length = self.n_elements

        if offset < 0 or length < 0 or (offset + length) > self.n_elements:
            raise IndexError(
                f"offset {offset} and length {length} do not describe a valid slice "
                f"of buffer with {self.n_elements} elements"
            )

        if _types.equal(self._data_type, _types.BOOL):
            return self._iter_bitmap(offset, length)
        else:
            return self._iter_dispatch(offset, length)

    def copy_into(self, dest, offset=0, length=None, dest_offset=0):
        if length is None:
            length = self.n_elements

        cdef Py_buffer buffer
        PyObject_GetBuffer(dest, &buffer, PyBUF_WRITABLE | PyBUF_ANY_CONTIGUOUS)

        cdef int64_t c_offset = offset
        cdef int64_t c_length = length
        cdef int64_t c_item_size = self.itemsize
        cdef int64_t c_dest_offset = dest_offset
        self._check_copy_into_bounds(&buffer, c_offset, c_length, dest_offset, c_item_size)

        cdef uint8_t* dest_uint8 = <uint8_t*>buffer.buf
        cdef int64_t dest_offset_bytes = c_dest_offset * c_item_size
        cdef int64_t src_offset_bytes = c_offset * c_item_size
        cdef int64_t bytes_to_copy = c_length * c_item_size

        memcpy(
            &(dest_uint8[dest_offset_bytes]),
            &(self._ptr.data.as_uint8[src_offset_bytes]),
            bytes_to_copy
        )

        PyBuffer_Release(&buffer)
        return bytes_to_copy

    def unpack_bits_into(self, dest, offset=0, length=None, dest_offset=0):
        if not _types.equal(self._data_type, _types.BOOL):
            raise ValueError("Can't unpack non-boolean buffer")

        if length is None:
            length = self.n_elements

        cdef Py_buffer buffer
        PyObject_GetBuffer(dest, &buffer, PyBUF_WRITABLE | PyBUF_ANY_CONTIGUOUS)
        self._check_copy_into_bounds(&buffer, offset, length, dest_offset, 1)

        ArrowBitsUnpackInt8(
            self._ptr.data.as_uint8,
            offset,
            length,
            &(<int8_t*>buffer.buf)[dest_offset]
        )

        PyBuffer_Release(&buffer)
        return length

    def unpack_bits(self, offset=0, length=None):
        if length is None:
            length = self.n_elements

        out = CBufferBuilder().set_format("?")
        out.reserve_bytes(length)
        self.unpack_bits_into(out, offset, length)
        out.advance(length)
        return out.finish()

    def copy(self, offset=0, length=None):
        if length is None:
            length = self.n_elements

        cdef int64_t bytes_to_copy = length * self.itemsize
        out = CBufferBuilder().set_data_type(self.data_type_id)
        out.reserve_bytes(bytes_to_copy)
        self.copy_into(out, offset, length)
        out.advance(bytes_to_copy)
        return out.finish()

    cdef _check_copy_into_bounds(self, Py_buffer* dest, int64_t offset, int64_t length,
                                 int64_t dest_offset, int64_t dest_itemsize):
        if offset < 0 or length < 0 or (offset + length) > self.n_elements:
            PyBuffer_Release(dest)
            raise IndexError(
                f"offset {offset} and length {length} do not describe a valid slice "
                f"of buffer with {self.n_elements} elements"
            )

        if dest.itemsize != 1 and dest.itemsize != dest_itemsize:
            raise ValueError(
                "Destination buffer must have itemsize == 1 or "
                f"itemsize == {dest_itemsize}"
            )

        cdef int64_t dest_offset_bytes = dest_offset * dest_itemsize
        cdef int64_t bytes_to_copy = dest_itemsize * length
        if dest_offset < 0 or dest.len < (dest_offset_bytes + bytes_to_copy):
            buffer_len = dest.len
            PyBuffer_Release(dest)
            raise IndexError(
                f"Can't unpack {length} elements into buffer of size {buffer_len} "
                f"with dest_offset = {dest_offset}"
            )

    def _iter_bitmap(self, int64_t offset, int64_t length):
        cdef uint8_t item
        cdef int64_t i

        if offset % 8 == 0:
            first_byte = offset // 8
            last_byte = self._shape
            i = 0

            for byte_i in range(first_byte, last_byte):
                item = self._ptr.data.as_uint8[byte_i]
                for j in range(8):
                    yield (item & (<uint8_t>1 << j)) != 0
                    i += 1
                    if i >= length:
                        return
        else:
            for i in range(length):
                yield ArrowBitGet(self._ptr.data.as_uint8, offset + i) != 0


    cdef Py_ssize_t _item_size(self):
        if self._element_size_bits < 8:
            return 1
        else:
            return self._element_size_bits // 8


    def __dlpack__(self, stream=None):
        """
        Export CBufferView as a DLPack capsule.

        Parameters
        ----------
        stream : int, optional
            A Python integer representing a pointer to a stream.
            Stream is provided by the consumer to the producer to instruct the producer
            to ensure that operations can safely be performed on the array.

        Returns
        -------
        capsule : PyCapsule
            A DLPack capsule for the array, pointing to a DLManagedTensor.
        """
        # Note: parent offset not applied!
        return view_to_dlpack(self, stream)


    def __dlpack_device__(self):
        """
        Return the DLPack device tuple this CBufferView resides on.

        Returns
        -------
        tuple : Tuple[int, int]
            Tuple with index specifying the type of the device (where
            CPU = 1, see python/src/nanoarrow/dpack_abi.h) and index of the
            device which is 0 by default for CPU.
        """
        cdef DLDevice dlpack_device = view_to_dlpack_device(self)
        return dlpack_device.device_type, dlpack_device.device_id

    # These are special methods, which can't be cdef and we can't
    # call them from elsewhere. We implement the logic for the buffer
    # protocol separately so we can re-use it in the CBuffer.
    def __getbuffer__(self, Py_buffer *buffer, int flags):
        self._do_getbuffer(buffer, flags)

    def __releasebuffer__(self, Py_buffer *buffer):
        self._do_releasebuffer(buffer)

    cdef _do_getbuffer(self, Py_buffer *buffer, int flags):
        if self.device is not DEVICE_CPU:
            raise RuntimeError("CBufferView is not a CPU buffer")

        if flags & PyBUF_WRITABLE:
            raise BufferError("CBufferView does not support PyBUF_WRITABLE")

        buffer.buf = <void*>self._ptr.data.data

        if flags & PyBUF_FORMAT:
            buffer.format = self._format
        else:
            buffer.format = NULL

        buffer.internal = NULL
        buffer.itemsize = self._strides
        buffer.len = self._ptr.size_bytes
        buffer.ndim = 1
        buffer.obj = self
        buffer.readonly = 1
        buffer.shape = &self._shape
        buffer.strides = &self._strides
        buffer.suboffsets = NULL

    cdef _do_releasebuffer(self, Py_buffer* buffer):
        pass

    def __repr__(self):
        class_label = _repr_utils.make_class_label(self, module="nanoarrow.c_buffer")
        return f"{class_label}({_repr_utils.buffer_view_repr(self)})"


cdef class CBuffer:
    """Wrapper around readable owned buffer content

    Like the CBufferView, the CBuffer represents readable buffer content; however,
    unlike the CBufferView, the CBuffer always represents a valid ArrowBuffer C object.
    Whereas the CBufferView is primarily concerned with accessing the contents of a
    buffer, the CBuffer is primarily concerned with managing ownership of an external
    buffer such that it exported as an Arrow array.
    """

    def __cinit__(self):
        self._base = None
        self._ptr = NULL
        self._data_type = <ArrowType>(_types.BINARY)
        self._element_size_bits = 0
        self._device = DEVICE_CPU
        # Set initial format to "B" (Cython makes this hard)
        self._format[0] = 66
        self._format[1] = 0
        self._get_buffer_count = 0
        self._view = CBufferView(
            None, 0,
            0, _types.BINARY, 0,
            CSharedSyncEvent(self._device)
        )

    cdef _assert_valid(self):
        if self._ptr == NULL:
            raise RuntimeError("CBuffer is not valid")

    cdef _assert_buffer_count_zero(self):
        if self._get_buffer_count != 0:
            raise RuntimeError(
                f"CBuffer already open ({self._get_buffer_count} ",
                f"references, {self._writable_get_buffer_count} writable)")

    cdef _populate_view(self):
        self._assert_valid()
        self._assert_buffer_count_zero()
        self._view = CBufferView(
            self._base, <uintptr_t>self._ptr.data,
            self._ptr.size_bytes, self._data_type, self._element_size_bits,
            CSharedSyncEvent(self._device)
        )

        snprintf(self._view._format, sizeof(self._view._format), "%s", self._format)

    def view(self):
        """Export this buffer as a CBufferView

        Returns a :class:`CBufferView` of this buffer. After calling this
        method, the original CBuffer will be invalidated and cannot be used.
        In general, the view of the buffer should be used to consume a buffer
        (whereas the CBuffer is primarily used to wrap an existing object in
        a way that it can be used to build a :class:`CArray`).
        """
        self._assert_valid()
        self._assert_buffer_count_zero()
        cdef ArrowBuffer* new_ptr
        self._view._base = _utils.alloc_c_buffer(&new_ptr)
        ArrowBufferMove(self._ptr, new_ptr)
        self._ptr = NULL
        return self._view

    @staticmethod
    def empty():
        """Create an empty CBuffer"""
        cdef CBuffer out = CBuffer()
        out._base = _utils.alloc_c_buffer(&out._ptr)
        return out

    @staticmethod
    def from_pybuffer(obj) -> CBuffer:
        """Create a CBuffer using the Python buffer protocol

        Wraps a buffer using the Python buffer protocol as a CBuffer that can be
        used to create an array.

        Parameters
        ----------
        obj : buffer-like
            The object on which to invoke the Python buffer protocol
        """
        cdef CBuffer out = CBuffer()
        out._base = _utils.alloc_c_buffer(&out._ptr)
        out._set_format(_utils.c_buffer_set_pybuffer(obj, &out._ptr))
        out._device = DEVICE_CPU
        out._populate_view()
        return out

    @staticmethod
    def from_dlpack(obj, stream=None) -> CBuffer:
        """Create a CBuffer using the DLPack protocol

        Wraps a tensor from an external library as a CBuffer that can be used
        to create an array.

        Parameters
        ----------
        obj : object with a ``__dlpack__`` attribute
            The object on which to invoke the DLPack protocol
        stream : int, optional
            The stream on which the tensor represented by obj should be made
            safe for use. This value is passed to the object's ``__dlpack__``
            method; however, the CBuffer does not keep any record of this (i.e.,
            the caller is responsible for creating a sync event after creating one
            or more buffers in this way).
        """
        capsule = obj.__dlpack__(stream=stream)
        cdef DLManagedTensor* dlm_tensor = <DLManagedTensor*>PyCapsule_GetPointer(
            capsule, "dltensor"
        )
        cdef DLTensor* dl_tensor = &dlm_tensor.dl_tensor

        if not dlpack_strides_are_contiguous(dl_tensor):
            raise ValueError("Non-contiguous dlpack strides not supported")

        cdef Device device = Device.resolve(
            dl_tensor.device.device_type,
            dl_tensor.device.device_id
        )
        cdef int arrow_type = dlpack_data_type_to_arrow(dl_tensor.dtype)
        cdef uint8_t* data_ptr = <uint8_t*>dl_tensor.data + dl_tensor.byte_offset

        cdef int64_t size_bytes = 1
        cdef int64_t element_size_bytes = dl_tensor.dtype.bits // 8
        for i in range(dl_tensor.ndim):
            size_bytes *= dl_tensor.shape[i] * element_size_bytes

        cdef CBuffer out = CBuffer()

        out._base = _utils.alloc_c_buffer(&out._ptr)
        _utils.c_buffer_set_pyobject(capsule, data_ptr, size_bytes, &out._ptr)
        out._set_data_type(arrow_type)
        out._device = device
        out._populate_view()

        return out

    def _set_format(self, str format):
        self._assert_buffer_count_zero()

        element_size_bytes, data_type = _types.from_format(format)
        self._data_type = data_type
        self._element_size_bits = element_size_bytes * 8
        format_bytes = format.encode("UTF-8")
        snprintf(self._format, sizeof(self._format), "%s", <const char*>format_bytes)
        self._populate_view()
        return self

    def _set_data_type(self, ArrowType type_id, int element_size_bits=0):
        self._assert_buffer_count_zero()

        self._element_size_bits = _types.to_format(
            type_id,
            element_size_bits,
            sizeof(self._format),
            self._format
        )
        self._data_type = type_id

        self._populate_view()
        return self

    def _addr(self):
        self._assert_valid()
        return <uintptr_t>self._ptr.data

    @property
    def _get_buffer_count(self):
        return self._get_buffer_count

    @property
    def size_bytes(self):
        self._assert_valid()
        return self._ptr.size_bytes

    @property
    def data_type(self):
        self._assert_valid()
        return ArrowTypeString(self._data_type).decode("UTF-8")

    @property
    def data_type_id(self):
        self._assert_valid()
        return self._data_type

    @property
    def element_size_bits(self):
        self._assert_valid()
        return self._element_size_bits

    @property
    def itemsize(self):
        self._assert_valid()
        return self._view.itemsize

    @property
    def format(self):
        self._assert_valid()
        return self._format.decode("UTF-8")

    @property
    def device(self):
        self._assert_valid()
        return self._view.device

    def __len__(self):
        self._assert_valid()
        return len(self._view)

    def __getitem__(self, k):
        self._assert_valid()
        return self._view[k]

    def __iter__(self):
        self._assert_valid()
        return iter(self._view)

    @property
    def n_elements(self):
        self._assert_valid()
        return self._view.n_elements

    def element(self, i):
        self._assert_valid()
        return self._view.element(i)

    def elements(self, offset=0, length=None):
        self._assert_valid()
        return self._view.elements(offset, length)

    def __getbuffer__(self, Py_buffer* buffer, int flags):
        self._assert_valid()
        self._view._do_getbuffer(buffer, flags)
        self._get_buffer_count += 1

    def __releasebuffer__(self, Py_buffer* buffer):
        if self._get_buffer_count <= 0:
            raise RuntimeError("CBuffer buffer reference count underflow (releasebuffer)")

        self._view._do_releasebuffer(buffer)
        self._get_buffer_count -= 1

    def __repr__(self):
        class_label = _repr_utils.make_class_label(self, module="nanoarrow.c_buffer")
        if self._ptr == NULL:
            return f"{class_label}(<invalid>)"

        return f"{class_label}({_repr_utils.buffer_view_repr(self._view)})"


cdef class CBufferBuilder:
    """Wrapper around writable CPU buffer content

    This class provides a growable type-aware buffer that can be used
    to create a typed buffer from a Python iterable. This method of
    creating a buffer is usually slower than constructors like
    ``array.array()`` or ``numpy.array()``; however, this class supports
    all Arrow types with a single data buffer (e.g., boolean bitmaps,
    float16, intervals, fixed-size binary), some of which are not supported
    by other constructors.
    """
    cdef CBuffer _buffer
    cdef bint _locked

    def __cinit__(self):
        self._buffer = CBuffer.empty()
        self._locked = False

    cdef _assert_unlocked(self):
        if self._locked:
            raise BufferError("CBufferBuilder is locked")

    # Implement the buffer protocol so that this object can be used as
    # the argument to Struct.readinto() (or perhaps written to by
    # an independent library).
    def __getbuffer__(self, Py_buffer* buffer, int flags):
        self._assert_unlocked()
        PyBuffer_FillInfo(
            buffer,
            self,
            self._buffer._ptr.data,
            self._buffer._ptr.capacity_bytes,
            0,
            flags
        )
        self._locked = True

    def __releasebuffer__(self, Py_buffer* buffer):
        self._locked = False

    def set_data_type(self, ArrowType type_id, int element_size_bits=0):
        """Set the data type used to interpret elements in :meth:`write_elements`."""
        self._buffer._set_data_type(type_id, element_size_bits)
        return self

    def set_format(self, str format):
        """Set the Python buffer format used to interpret elements in
        :meth:`write_elements`.
        """
        self._buffer._set_format(format)
        return self

    @property
    def format(self):
        """The ``struct`` format code of the underlying buffer"""
        return self._buffer._format.decode()

    @property
    def size_bytes(self):
        """The number of bytes that have been written to this buffer"""
        return self._buffer.size_bytes

    @property
    def itemsize(self):
        return self._buffer.itemsize

    def __len__(self):
        return self._buffer.size_bytes // self.itemsize

    @property
    def capacity_bytes(self):
        """The number of bytes allocated in the underlying buffer"""
        return self._buffer._ptr.capacity_bytes

    def reserve_bytes(self, int64_t additional_bytes):
        """Ensure that the underlying buffer has space for ``additional_bytes``
        more bytes to be written"""
        self._assert_unlocked()
        cdef int code = ArrowBufferReserve(self._buffer._ptr, additional_bytes)
        _utils.Error.raise_error_not_ok("ArrowBufferReserve()", code)
        return self

    def advance(self, int64_t additional_bytes):
        """Manually increase :attr:`size_bytes` by ``additional_bytes``

        This can be used after writing to the buffer using the buffer protocol
        to ensure that :attr:`size_bytes` accurately reflects the number of
        bytes written to the buffer.
        """
        cdef int64_t new_size = self._buffer._ptr.size_bytes + additional_bytes
        if new_size < 0 or new_size > self._buffer._ptr.capacity_bytes:
            raise IndexError(f"Can't advance {additional_bytes} from {self.size_bytes}")

        self._buffer._ptr.size_bytes = new_size
        return self

    def write_fill(self, uint8_t value, int64_t size_bytes):
        """Write fill bytes to this buffer

        Appends the byte ``value`` to this buffer ``size_bytes`` times.
        """
        self._assert_unlocked()
        cdef int code = ArrowBufferAppendFill(self._buffer._ptr, value, size_bytes)
        _utils.Error.raise_error_not_ok("ArrowBufferAppendFill", code)

    def write(self, content):
        """Write bytes to this buffer

        Writes the bytes of ``content`` without considering the element type of
        ``content`` or the element type of this buffer.

        This method returns the number of bytes that were written.
        """
        self._assert_unlocked()

        cdef Py_buffer buffer
        cdef int64_t out
        PyObject_GetBuffer(content, &buffer, PyBUF_ANY_CONTIGUOUS)

        # TODO: Check for single dimension?

        cdef int code = ArrowBufferReserve(self._buffer._ptr, buffer.len)
        if code != NANOARROW_OK:
            PyBuffer_Release(&buffer)
            _utils.Error.raise_error("ArrowBufferReserve()", code)

        code = PyBuffer_ToContiguous(
            self._buffer._ptr.data + self._buffer._ptr.size_bytes,
            &buffer,
            buffer.len,
            # 'C' (not sure how to pass a character literal here)
            43
        )
        out = buffer.len
        PyBuffer_Release(&buffer)
        _utils.Error.raise_error_not_ok("PyBuffer_ToContiguous()", code)

        self._buffer._ptr.size_bytes += out
        return out

    def write_elements(self, obj):
        """"Write an iterable of elements to this buffer

        Writes the elements of iterable ``obj`` according to the binary
        representation specified by :attr:`format`. This is currently
        powered by ``struct.pack_into()`` except when building bitmaps
        where an internal implementation is used.

        This method returns the number of elements that were written.
        """
        self._assert_unlocked()

        # Boolean arrays need their own writer since Python provides
        # no helpers for bitpacking
        if _types.equal(self._buffer._data_type, _types.BOOL):
            return self._write_bits(obj)

        cdef int64_t n_values = 0
        cdef int64_t bytes_per_element = calcsize(self._buffer.format)
        cdef int code

        struct_obj = Struct(self._buffer._format)
        pack = struct_obj.pack
        write = self.write

        # If an object has a length, we can avoid extra allocations
        if hasattr(obj, "__len__"):
            code = ArrowBufferReserve(self._buffer._ptr, bytes_per_element * len(obj))
            _utils.Error.raise_error_not_ok("ArrowBufferReserve()", code)

        # Types whose Python representation is a tuple need a slightly different
        # pack_into() call. if code != NANOARROW_OK is used instead of
        # Error.raise_error_not_ok() Cython can avoid the extra function call
        # and this is a very tight loop.
        if _types.one_of(
            self._buffer._data_type,
            (_types.INTERVAL_DAY_TIME, _types.INTERVAL_MONTH_DAY_NANO)
        ):
            for item in obj:
                write(pack(*item))
                n_values += 1

        else:
            for item in obj:
                write(pack(item))
                n_values += 1

        return n_values

    cdef _write_bits(self, obj):
        if self._buffer._ptr.size_bytes != 0:
            raise NotImplementedError("Append to bitmap that has already been appended to")

        cdef char buffer_item = 0
        cdef int buffer_item_i = 0
        cdef int code
        cdef int64_t n_values = 0
        for item in obj:
            n_values += 1
            if item:
                buffer_item |= (<char>1 << buffer_item_i)

            buffer_item_i += 1
            if buffer_item_i == 8:
                code = ArrowBufferAppendInt8(self._buffer._ptr, buffer_item)
                _utils.Error.raise_error_not_ok("ArrowBufferAppendInt8()", code)
                buffer_item = 0
                buffer_item_i = 0

        if buffer_item_i != 0:
            code = ArrowBufferAppendInt8(self._buffer._ptr, buffer_item)
            _utils.Error.raise_error_not_ok("ArrowBufferAppendInt8()", code)

        return n_values

    def finish(self):
        """Finish building this buffer

        Performs any steps required to finish building this buffer and
        returns the result. Any behaviour resulting from calling methods
        on this object after it has been finished is not currently
        defined (but should not crash).
        """
        self._assert_unlocked()

        cdef CBuffer out = self._buffer
        out._populate_view()

        self._buffer = CBuffer.empty()
        return out

    def __repr__(self):
        class_label = _repr_utils.make_class_label(self, module="nanoarrow.c_buffer")
        return f"{class_label}({self.size_bytes}/{self.capacity_bytes})"


cdef class NoneAwareWrapperIterator:
    """Nullable iterator wrapper

    This class wraps an iterable ``obj`` that might contain ``None`` values
    such that the iterable provided by this class contains "empty" (but valid)
    values. After ``obj`` has been completely consumed, one can call
    ``finish()`` to obtain the resulting bitmap. This is useful for passing
    iterables that might contain None to tools that cannot handle them
    (e.g., struct.pack(), array.array()).
    """
    cdef ArrowBitmap _bitmap
    cdef object _obj
    cdef object _value_if_none
    cdef int64_t _valid_count
    cdef int64_t _item_count

    def __cinit__(self, obj, type_id, item_size_bytes=0):
        ArrowBitmapInit(&self._bitmap)
        self._obj = iter(obj)

        self._value_if_none = self._get_value_if_none(type_id, item_size_bytes)
        self._valid_count = 0
        self._item_count = 0

    def __dealloc__(self):
        ArrowBitmapReset(&self._bitmap)

    def reserve(self, int64_t additional_elements):
        cdef int code = ArrowBitmapReserve(&self._bitmap, additional_elements)
        _utils.Error.raise_error_not_ok(self, code)

    def _get_value_if_none(self, type_id, item_size_bytes=0):
        if _types.equal(type_id, _types.INTERVAL_DAY_TIME):
            return (0, 0)
        elif _types.equal(type_id, _types.INTERVAL_MONTH_DAY_NANO):
            return (0, 0, 0)
        elif _types.equal(type_id, _types.BOOL):
            return False
        elif _types.one_of(type_id, (_types.BINARY, _types.FIXED_SIZE_BINARY)):
            return b"\x00" * item_size_bytes
        elif _types.one_of(type_id, (_types.HALF_FLOAT, _types.FLOAT, _types.DOUBLE)):
            return 0.0
        else:
            return 0

    cdef _append_to_validity(self, int is_valid):
        self._valid_count += is_valid
        self._item_count += 1

        # Avoid allocating a bitmap if all values seen so far are valid
        if self._valid_count == self._item_count:
            return

        # If the bitmap hasn't been allocated yet, allocate it now and
        # fill with 1s for all previous elements.
        cdef int code
        if self._bitmap.size_bits == 0 and self._item_count > 1:
            code = ArrowBitmapAppend(&self._bitmap, 1, self._item_count - 1)
            if code != NANOARROW_OK:
                _utils.Error.raise_error("ArrowBitmapAppend()", code)

        # Append this element to the bitmap
        code = ArrowBitmapAppend(&self._bitmap, is_valid, 1)
        if code != NANOARROW_OK:
            _utils.Error.raise_error("ArrowBitmapAppend()", code)

    def __iter__(self):
        for item in self._obj:
            if item is None:
                self._append_to_validity(0)
                yield self._value_if_none
            else:
                self._append_to_validity(1)
                yield item

    def finish(self):
        """Obtain the total count, null count, and validity bitmap after
        consuming this iterable."""
        cdef CBuffer validity
        null_count = self._item_count - self._valid_count

        if null_count > 0:
            validity = CBuffer.empty()
            ArrowBufferMove(&self._bitmap.buffer, validity._ptr)
            validity._set_data_type(<ArrowType>_types.BOOL)

        return (
            self._item_count,
            null_count,
            validity if null_count > 0 else None
        )