# 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 int32_t, int64_t, uintptr_t
from cpython.bytes cimport PyBytes_FromStringAndSize, PyBytes_AsString, PyBytes_Size
from cpython.pycapsule cimport PyCapsule_GetPointer

from nanoarrow_c cimport (
    ARROW_FLAG_DICTIONARY_ORDERED,
    ARROW_FLAG_MAP_KEYS_SORTED,
    ARROW_FLAG_NULLABLE,
    ArrowFree,
    ArrowLayout,
    ArrowMalloc,
    ArrowMetadataBuilderAppend,
    ArrowMetadataBuilderInit,
    ArrowMetadataReaderInit,
    ArrowMetadataReaderRead,
    ArrowSchema,
    ArrowSchemaAllocateChildren,
    ArrowSchemaAllocateDictionary,
    ArrowSchemaDeepCopy,
    ArrowSchemaInit,
    ArrowSchemaMove,
    ArrowSchemaRelease,
    ArrowSchemaSetMetadata,
    ArrowSchemaSetType,
    ArrowSchemaSetTypeDateTime,
    ArrowSchemaSetTypeDecimal,
    ArrowSchemaSetTypeFixedSize,
    ArrowSchemaSetFormat,
    ArrowSchemaSetName,
    ArrowSchemaToString,
    ArrowSchemaViewInit,
    ArrowStringView,
    ArrowTimeUnit,
    ArrowTimeUnitString,
    ArrowType,
    ArrowTypeString,
    NANOARROW_BUFFER_TYPE_NONE,
    NANOARROW_MAX_FIXED_BUFFERS,
    NANOARROW_TIME_UNIT_SECOND,
    NANOARROW_TIME_UNIT_MILLI,
    NANOARROW_TIME_UNIT_MICRO,
    NANOARROW_TIME_UNIT_NANO,
)

from nanoarrow cimport _types
from nanoarrow._buffer cimport CBuffer
from nanoarrow._utils cimport alloc_c_schema, Error

from typing import Iterable, List, Mapping, Tuple, Union

from nanoarrow import _repr_utils


# This is likely a better fit for a dedicated testing module; however, we need
# it here to produce nice error messages when ensuring that one or
# more arrays conform to a given or inferred schema.
cpdef assert_type_equal(actual, expected, bint check_nullability):
    """Test two schemas for data type equality

    Checks two CSchema objects for type equality (i.e., that an array with
    schema ``actual`` contains elements with the same logical meaning as and
    array with schema ``expected``). Notably, this excludes metadata from
    all nodes in the schema.

    Parameters
    ----------
    actual : CSchema
        The schema to be tested for equality
    expected : CSchema
        The schema against which to test
    check_nullability : bool
        If True, actual and expected will be considered equal if their
        data type information and marked nullability are identical.
    """
    if not isinstance(actual, CSchema):
        raise TypeError(f"actual is {type(actual).__name__}, not CSchema")

    if not isinstance(expected, CSchema):
        raise TypeError(f"expected is {type(expected).__name__}, not CSchema")

    if not actual.type_equals(expected, check_nullability=check_nullability):
        actual_label = actual._to_string(max_chars=80, recursive=True)
        expected_label = expected._to_string(max_chars=80, recursive=True)
        raise ValueError(
            f"Expected schema\n  '{expected_label}'"
            f"\nbut got\n  '{actual_label}'"
        )


cdef class CArrowTimeUnit:
    """
    Wrapper around ArrowTimeUnit to provide implementations in Python access
    to the values.
    """

    SECOND = NANOARROW_TIME_UNIT_SECOND
    MILLI = NANOARROW_TIME_UNIT_MILLI
    MICRO = NANOARROW_TIME_UNIT_MICRO
    NANO = NANOARROW_TIME_UNIT_NANO


cdef class CLayout:
    """Abstract buffer information for Arrow types

    Provides accessors for buffer counts, types, and attributes.
    """

    def __cinit__(self, base, uintptr_t ptr):
        self._base = base
        self._layout = <ArrowLayout*>ptr

        self._n_buffers = NANOARROW_MAX_FIXED_BUFFERS
        for i in range(NANOARROW_MAX_FIXED_BUFFERS):
            if self._layout.buffer_type[i] == NANOARROW_BUFFER_TYPE_NONE:
                self._n_buffers = i
                break

    @property
    def n_buffers(self) -> int:
        return self._n_buffers

    @property
    def buffer_data_type_id(self) -> int:
        return tuple(self._layout.buffer_data_type[i] for i in range(self._n_buffers))

    @property
    def element_size_bits(self) -> int:
        return tuple(self._layout.element_size_bits[i] for i in range(self._n_buffers))

    @property
    def child_size_elements(self) -> int:
        return self._layout.child_size_elements


cdef class SchemaMetadata:
    """Dictionary-like wrapper around a lazily-parsed CSchema.metadata string

    The Arrow C Data interface encodes key/value metadata as a bytes-to-bytes
    mapping using a specific packed binary encoding. This class maintains a
    reference to the underlying storage and parses it as required. Note that
    unlike a Python dictionary, ``SchemaMetadata`` can contain duplicate
    keys.
    """

    def __cinit__(self, object base, uintptr_t ptr):
        self._base = base
        self._metadata = <const char*>ptr

    @staticmethod
    def empty():
        """Create an empty SchemaMetadata with no keys or values"""
        return SchemaMetadata(None, 0)

    cdef _init_reader(self):
        cdef int code = ArrowMetadataReaderInit(&self._reader, self._metadata)
        Error.raise_error_not_ok("ArrowMetadataReaderInit()", code)

    def __len__(self):
        self._init_reader()
        return self._reader.remaining_keys

    def __contains__(self, item):
        for key, _ in self.items():
            if item == key:
                return True

        return False

    def __getitem__(self, k) -> bytes:
        """Get the value associated with a unique key

        Retrieves the unique value associated with k. Raises KeyError if
        k does not point to exactly one value in the metadata.
        """
        out = None

        for key, value in self.items():
            if k == key:
                if out is None:
                    out = value
                else:
                    raise KeyError(f"key {k} matches more than one value in metadata")

        if out is None:
            raise KeyError(f"Key {k} not found")

        return out

    def __iter__(self):
        for key, _ in self.items():
            yield key

    def keys(self) -> List[bytes]:
        """List meadata keys

        The result may contain duplicate keys if they exist in the metadata.
        """
        return list(self)

    def values(self) -> List[bytes]:
        """List metadata values"""
        return [value for _, value in self.items()]

    def items(self) -> Iterable[bytes, bytes]:
        """Iterate over key/value pairs

        The result may contain duplicate keys if they exist in the metadata."""
        cdef ArrowStringView key
        cdef ArrowStringView value
        self._init_reader()
        while self._reader.remaining_keys > 0:
            ArrowMetadataReaderRead(&self._reader, &key, &value)
            key_obj = PyBytes_FromStringAndSize(key.data, key.size_bytes)
            value_obj = PyBytes_FromStringAndSize(value.data, value.size_bytes)
            yield key_obj, value_obj

    def __repr__(self) -> str:
        lines = [
            f"<{_repr_utils.make_class_label(self)}>",
            _repr_utils.metadata_repr(self)
        ]
        return "\n".join(lines)


cdef class CSchema:
    """Low-level ArrowSchema wrapper

    This object is a literal wrapper around a read-only ArrowSchema. It provides field accessors
    that return Python objects and handles the C Data interface lifecycle (i.e., initialized
    ArrowSchema structures are always released).

    See ``nanoarrow.c_schema()`` for construction and usage examples.
    """

    @staticmethod
    def allocate() -> CSchema:
        """Allocate a released CSchema"""
        cdef ArrowSchema* c_schema_out
        base = alloc_c_schema(&c_schema_out)
        return CSchema(base, <uintptr_t>(c_schema_out))

    def __cinit__(self, object base, uintptr_t addr):
        self._base = base
        self._ptr = <ArrowSchema*>addr

    def __deepcopy__(self, memo=None) -> CSchema:
        cdef CSchema out = CSchema.allocate()
        cdef int code = ArrowSchemaDeepCopy(self._ptr, out._ptr)
        Error.raise_error_not_ok("ArrowSchemaDeepCopy()", code)

        return out

    @staticmethod
    def _import_from_c_capsule(schema_capsule) -> CSchema:
        """Import from a ArrowSchema PyCapsule

        Parameters
        ----------
        schema_capsule : PyCapsule
            A valid PyCapsule with name 'arrow_schema' containing an
            ArrowSchema pointer.
        """
        return CSchema(
            schema_capsule,
            <uintptr_t>PyCapsule_GetPointer(schema_capsule, "arrow_schema")
        )

    def __arrow_c_schema__(self):
        """
        Export to a ArrowSchema PyCapsule
        """
        self._assert_valid()

        cdef ArrowSchema* c_schema_out
        schema_capsule = alloc_c_schema(&c_schema_out)

        cdef int code = ArrowSchemaDeepCopy(self._ptr, c_schema_out)
        Error.raise_error_not_ok("ArrowSchemaDeepCopy", code)
        return schema_capsule

    @property
    def _capsule(self):
        """
        Returns the capsule backing this CSchema or None if it does not exist
        or points to a parent ArrowSchema.
        """
        cdef ArrowSchema* maybe_capsule_ptr
        maybe_capsule_ptr = <ArrowSchema*>PyCapsule_GetPointer(self._base, 'arrow_schema')

        # This will return False if this is a child CSchema whose capsule holds
        # the parent ArrowSchema
        if maybe_capsule_ptr == self._ptr:
            return self._base

        return None

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

    def is_valid(self) -> bool:
        """Check for a non-null and non-released underlying ArrowSchema"""
        return self._ptr != NULL and self._ptr.release != NULL

    def _assert_valid(self):
        if self._ptr == NULL:
            raise RuntimeError("schema is NULL")
        if self._ptr.release == NULL:
            raise RuntimeError("schema is released")

    def _to_string(self, int64_t max_chars=0, recursive=False) -> str:
        cdef int64_t n_chars
        if max_chars == 0:
            n_chars = ArrowSchemaToString(self._ptr, NULL, 0, recursive)
        else:
            n_chars = max_chars

        cdef char* out = <char*>ArrowMalloc(n_chars + 1)
        if not out:
            raise MemoryError()

        ArrowSchemaToString(self._ptr, out, n_chars + 1, recursive)
        out_str = out.decode("UTF-8")
        ArrowFree(out)

        return out_str

    def __repr__(self) -> str:
        return _repr_utils.schema_repr(self)

    def type_equals(self, CSchema other, check_nullability: bool=False) -> bool:
        """Test two schemas for data type equality

        Checks two CSchema objects for type equality (i.e., that an array with
        schema ``actual`` contains elements with the same logical meaning as and
        array with schema ``expected``). Notably, this excludes metadata from
        all nodes in the schema.

        Parameters
        ----------
        other : CSchema
            The schema against which to test
        check_nullability : bool
            If True, actual and expected will be considered equal if their
            data type information and marked nullability are identical.
        """
        self._assert_valid()

        if self._ptr == other._ptr:
            return True

        if self.format != other.format:
            return False

        # Nullability is not strictly part of the "type"; however, performing
        # this check recursively is verbose to otherwise accomplish and
        # sometimes this does matter.
        cdef int64_t flags = self.flags
        cdef int64_t other_flags = other.flags
        if not check_nullability:
            flags &= ~ARROW_FLAG_NULLABLE
            other_flags &= ~ARROW_FLAG_NULLABLE

        if flags != other_flags:
            return False

        if self.n_children != other.n_children:
            return False

        for child, other_child in zip(self.children, other.children):
            if not child.type_equals(other_child, check_nullability=check_nullability):
                return False

        if (self.dictionary is None) != (other.dictionary is None):
            return False

        if self.dictionary is not None:
            if not self.dictionary.type_equals(
                other.dictionary,
                check_nullability=check_nullability
            ):
                return False

        return True


    @property
    def format(self) -> str:
        self._assert_valid()
        if self._ptr.format != NULL:
            return self._ptr.format.decode()

    @property
    def name(self) -> Union[str, None]:
        self._assert_valid()
        if self._ptr.name != NULL:
            return self._ptr.name.decode()
        else:
            return None

    @property
    def flags(self) -> int:
        return self._ptr.flags

    @property
    def metadata(self) -> SchemaMetadata:
        self._assert_valid()
        if self._ptr.metadata != NULL:
            return SchemaMetadata(self._base, <uintptr_t>self._ptr.metadata)
        else:
            return None

    @property
    def n_children(self) -> int:
        self._assert_valid()
        return self._ptr.n_children

    def child(self, int64_t i):
        self._assert_valid()
        if i < 0 or i >= self._ptr.n_children:
            raise IndexError(f"{i} out of range [0, {self._ptr.n_children})")

        return CSchema(self._base, <uintptr_t>self._ptr.children[i])

    @property
    def children(self) -> Iterable[CSchema]:
        for i in range(self.n_children):
            yield self.child(i)

    @property
    def dictionary(self) -> Union[CSchema, None]:
        self._assert_valid()
        if self._ptr.dictionary != NULL:
            return CSchema(self, <uintptr_t>self._ptr.dictionary)
        else:
            return None

    def modify(self, *, format=None, name=None, flags=None, nullable=None,
               metadata=None, children=None, dictionary=None, validate=True) -> CSchema:
        cdef CSchemaBuilder builder = CSchemaBuilder.allocate()

        if format is None:
            builder.set_format(self.format)
        else:
            builder.set_format(format)

        if name is None:
            builder.set_name(self.name)
        elif name is not False:
            builder.set_name(name)

        if flags is None:
            builder.set_flags(self.flags)
        else:
            builder.set_flags(flags)

        if nullable is not None:
            builder.set_nullable(nullable)

        if metadata is None:
            if self.metadata is not None:
                builder.append_metadata(self.metadata)
        else:
            builder.append_metadata(metadata)

        if children is None:
            if self.n_children > 0:
                builder.allocate_children(self.n_children)
                for i, child in enumerate(self.children):
                    builder.set_child(i, None, child)
        elif hasattr(children, "items"):
            builder.allocate_children(len(children))
            for i, item in enumerate(children.items()):
                name, child = item
                builder.set_child(i, name, child)
        else:
            builder.allocate_children(len(children))
            for i, child in enumerate(children):
                builder.set_child(i, None, child)

        if dictionary is None:
            if self.dictionary:
                builder.set_dictionary(self.dictionary)
        elif dictionary is not False:
            builder.set_dictionary(dictionary)

        if validate:
            builder.validate()

        return builder.finish()


cdef class CSchemaView:
    """Low-level ArrowSchemaView wrapper

    This object is a literal wrapper around a read-only ArrowSchemaView. It provides field accessors
    that return Python objects and handles structure lifecycle. Compared to an ArrowSchema,
    the nanoarrow ArrowSchemaView facilitates access to the deserialized content of an ArrowSchema
    (e.g., parameter values for parameterized types).

    See `nanoarrow.c_schema_view()` for construction and usage examples.
    """

    def __cinit__(self, CSchema schema):
        self._base = schema
        self._schema_view.type = <ArrowType>_types.UNINITIALIZED
        self._schema_view.storage_type = <ArrowType>_types.UNINITIALIZED

        cdef Error error = Error()
        cdef int code = ArrowSchemaViewInit(&self._schema_view, schema._ptr, &error.c_error)
        error.raise_message_not_ok("ArrowSchemaViewInit()", code)

        self._dictionary_ordered = schema._ptr.flags & ARROW_FLAG_DICTIONARY_ORDERED
        self._nullable = schema._ptr.flags & ARROW_FLAG_NULLABLE
        self._map_keys_sorted = schema._ptr.flags & ARROW_FLAG_MAP_KEYS_SORTED

    @property
    def layout(self) -> CLayout:
        return CLayout(self, <uintptr_t>&self._schema_view.layout)

    @property
    def type_id(self) -> int:
        return self._schema_view.type

    @property
    def storage_type_id(self) -> int:
        return self._schema_view.storage_type

    @property
    def storage_buffer_format(self) -> Union[str, None]:
        if self.buffer_format is not None:
            return self.buffer_format
        elif _types.equal(self._schema_view.type, _types.DATE32):
            return 'i'
        elif _types.one_of(
            self._schema_view.type,
            (_types.TIMESTAMP, _types.DATE64, _types.DURATION)
        ):
            return 'q'
        elif self.extension_name:
            return self._get_buffer_format()
        else:
            return None

    @property
    def buffer_format(self) -> Union[str, None]:
        """The Python struct format representing an element of this type
        or None if there is no Python format string that can represent this
        type without loosing information.
        """
        if self.extension_name:
            return None
        else:
            return self._get_buffer_format()

    def _get_buffer_format(self):
        if self._schema_view.type != self._schema_view.storage_type:
            return None

        # String/binary types do not have format strings as far as the Python
        # buffer protocol is concerned
        if self.layout.n_buffers != 2:
            return None

        cdef char out[128]
        cdef int element_size_bits = 0
        if _types.equal(self._schema_view.type, _types.FIXED_SIZE_BINARY):
            element_size_bits = self._schema_view.fixed_size * 8

        try:
            _types.to_format(self._schema_view.type, element_size_bits, sizeof(out), out)
            return out.decode()
        except ValueError:
            return None

    @property
    def type(self) -> str:
        cdef const char* type_str = ArrowTypeString(self._schema_view.type)
        if type_str != NULL:
            return type_str.decode()
        else:
            raise ValueError("ArrowTypeString() returned NULL")

    @property
    def storage_type(self) -> str:
        cdef const char* type_str = ArrowTypeString(self._schema_view.storage_type)
        if type_str != NULL:
            return type_str.decode()
        else:
            raise ValueError("ArrowTypeString() returned NULL")

    @property
    def dictionary_ordered(self) -> Union[bool, None]:
        if _types.equal(self._schema_view.type, _types.DICTIONARY):
            return self._dictionary_ordered != 0
        else:
            return None

    @property
    def nullable(self) -> bool:
        return self._nullable != 0

    @property
    def map_keys_sorted(self) -> Union[bool, None]:
        if _types.equal(self._schema_view.type, _types.MAP):
            return self._map_keys_sorted != 0
        else:
            return None

    @property
    def fixed_size(self) -> Union[bool, None]:
        if _types.is_fixed_size(self._schema_view.type):
            return self._schema_view.fixed_size
        else:
            return None

    @property
    def decimal_bitwidth(self) -> Union[int, None]:
        if _types.is_decimal(self._schema_view.type):
            return self._schema_view.decimal_bitwidth
        else:
            return None

    @property
    def decimal_precision(self) -> Union[int, None]:
        if _types.is_decimal(self._schema_view.type):
            return self._schema_view.decimal_precision
        else:
            return None

    @property
    def decimal_scale(self) -> Union[int, None]:
        if _types.is_decimal(self._schema_view.type):
            return self._schema_view.decimal_scale
        else:
            return None

    @property
    def time_unit_id(self) -> Union[int, None]:
        if _types.has_time_unit(self._schema_view.type):
            return self._schema_view.time_unit
        else:
            return None

    @property
    def time_unit(self) -> Union[str, None]:
        if _types.has_time_unit(self._schema_view.type):
            return ArrowTimeUnitString(self._schema_view.time_unit).decode()
        else:
            return None

    @property
    def timezone(self) -> Union[str, None]:
        if _types.equal(self._schema_view.type, _types.TIMESTAMP):
            return self._schema_view.timezone.decode()
        else:
            return None

    @property
    def union_type_ids(self) -> Union[Tuple[int, ...], None]:
        if _types.is_union(self._schema_view.type):
            type_ids_str = self._schema_view.union_type_ids.decode().split(',')
            return (int(type_id) for type_id in type_ids_str)
        else:
            return None

    @property
    def extension_name(self) -> Union[str, None]:
        if self._schema_view.extension_name.data != NULL:
            name_bytes = PyBytes_FromStringAndSize(
                self._schema_view.extension_name.data,
                self._schema_view.extension_name.size_bytes
            )
            return name_bytes.decode()
        else:
            return None

    @property
    def extension_metadata(self) -> Union[bytes, None]:
        if self._schema_view.extension_name.data != NULL:
            return PyBytes_FromStringAndSize(
                self._schema_view.extension_metadata.data,
                self._schema_view.extension_metadata.size_bytes
            )
        else:
            return None

    def __repr__(self) -> str:
        return _repr_utils.schema_view_repr(self)


cdef class CSchemaBuilder:
    """Helper for constructing an ArrowSchema

    The primary function of this class is to wrap the nanoarrow C library calls
    that build up the components of an ArrowSchema.
    """

    def __cinit__(self, CSchema schema):
        self.c_schema = schema
        self._ptr = schema._ptr
        if self._ptr.release == NULL:
            ArrowSchemaInit(self._ptr)

    @staticmethod
    def allocate() -> CSchemaBuilder:
        """Create a CSchemaBuilder

        Allocates memory for an ArrowSchema and populates it with nanoarrow's
        ArrowSchema private_data/release callback implementation. This should
        usually be followed by :meth:`set_type` or :meth:`set_format`.
        """
        return CSchemaBuilder(CSchema.allocate())

    def append_metadata(self, metadata: Mapping[bytes, bytes]) -> CSchemaBuilder:
        """Append key/value metadata"""
        cdef CBuffer buffer = CBuffer.empty()

        cdef const char* existing_metadata = self.c_schema._ptr.metadata
        cdef int code = ArrowMetadataBuilderInit(buffer._ptr, existing_metadata)
        Error.raise_error_not_ok("ArrowMetadataBuilderInit()", code)

        cdef ArrowStringView key
        cdef ArrowStringView value
        cdef int32_t keys_added = 0

        for k, v in metadata.items():
            k = k.encode() if isinstance(k, str) else bytes(k)
            key.data = PyBytes_AsString(k)
            key.size_bytes = PyBytes_Size(k)

            v = v.encode() if isinstance(v, str) else bytes(v)
            value.data = PyBytes_AsString(v)
            value.size_bytes = PyBytes_Size(v)

            code = ArrowMetadataBuilderAppend(buffer._ptr, key, value)
            Error.raise_error_not_ok("ArrowMetadataBuilderAppend()", code)

            keys_added += 1

        if keys_added > 0:
            code = ArrowSchemaSetMetadata(self.c_schema._ptr, <const char*>buffer._ptr.data)
            Error.raise_error_not_ok("ArrowSchemaSetMetadata()", code)

        return self

    def child(self, int64_t i) -> CSchemaBuilder:
        return CSchemaBuilder(self.c_schema.child(i))

    def set_type(self, int type_id) -> CSchemaBuilder:
        self.c_schema._assert_valid()

        cdef int code = ArrowSchemaSetType(self._ptr, <ArrowType>type_id)
        Error.raise_error_not_ok("ArrowSchemaSetType()", code)

        return self

    def set_type_decimal(self, int type_id, int precision, int scale) -> CSchemaBuilder:
        self.c_schema._assert_valid()

        cdef int code = ArrowSchemaSetTypeDecimal(self._ptr, <ArrowType>type_id, precision, scale)
        Error.raise_error_not_ok("ArrowSchemaSetType()", code)

    def set_type_fixed_size(self, int type_id, int fixed_size) -> CSchemaBuilder:
        self.c_schema._assert_valid()

        cdef int code = ArrowSchemaSetTypeFixedSize(self._ptr, <ArrowType>type_id, fixed_size)
        Error.raise_error_not_ok("ArrowSchemaSetTypeFixedSize()", code)

        return self

    def set_type_date_time(self, int type_id, int time_unit, timezone) -> CSchemaBuilder:
        self.c_schema._assert_valid()

        cdef int code
        if timezone is None:
            code = ArrowSchemaSetTypeDateTime(self._ptr, <ArrowType>type_id, <ArrowTimeUnit>time_unit, NULL)
        else:
            timezone = str(timezone)
            code = ArrowSchemaSetTypeDateTime(self._ptr, <ArrowType>type_id, <ArrowTimeUnit>time_unit, timezone.encode("UTF-8"))

        Error.raise_error_not_ok("ArrowSchemaSetTypeDateTime()", code)

        return self

    def set_format(self, str format) -> CSchemaBuilder:
        self.c_schema._assert_valid()

        cdef int code = ArrowSchemaSetFormat(self._ptr, format.encode("UTF-8"))
        Error.raise_error_not_ok("ArrowSchemaSetFormat()", code)

        return self

    def set_name(self, name) -> CSchemaBuilder:
        self.c_schema._assert_valid()

        cdef int code
        if name is None:
            code = ArrowSchemaSetName(self._ptr, NULL)
        else:
            name = str(name)
            code = ArrowSchemaSetName(self._ptr, name.encode("UTF-8"))

        Error.raise_error_not_ok("ArrowSchemaSetName()", code)

        return self

    def allocate_children(self, int n) -> CSchemaBuilder:
        self.c_schema._assert_valid()

        cdef int code = ArrowSchemaAllocateChildren(self._ptr, n)
        Error.raise_error_not_ok("ArrowSchemaAllocateChildren()", code)

        return self

    def set_child(self, int64_t i, name, CSchema child_src) -> CSchemaBuilder:
        self.c_schema._assert_valid()

        if i < 0 or i >= self._ptr.n_children:
            raise IndexError(f"Index out of range: {i}")

        if self._ptr.children[i].release != NULL:
            ArrowSchemaRelease(self._ptr.children[i])

        cdef int code = ArrowSchemaDeepCopy(child_src._ptr, self._ptr.children[i])
        Error.raise_error_not_ok("ArrowSchemaDeepCopy()", code)

        if name is not None:
            name = str(name)
            code = ArrowSchemaSetName(self._ptr.children[i], name.encode("UTF-8"))
            Error.raise_error_not_ok("ArrowSchemaSetName()", code)

        return self

    def set_dictionary(self, CSchema dictionary) -> CSchemaBuilder:
        self.c_schema._assert_valid()

        cdef int code
        if self._ptr.dictionary == NULL:
            code = ArrowSchemaAllocateDictionary(self._ptr)
            Error.raise_error_not_ok("ArrowSchemaAllocateDictionary()", code)

        if self._ptr.dictionary.release != NULL:
            ArrowSchemaRelease(self._ptr.dictionary)

        code = ArrowSchemaDeepCopy(dictionary._ptr, self._ptr.dictionary)
        Error.raise_error_not_ok("ArrowSchemaDeepCopy()", code)

        return self

    def set_flags(self, flags) -> CSchemaBuilder:
        self._ptr.flags = flags
        return self

    def set_nullable(self, nullable) -> CSchemaBuilder:
        if nullable:
            self._ptr.flags = self._ptr.flags | ARROW_FLAG_NULLABLE
        else:
            self._ptr.flags = self._ptr.flags & ~ARROW_FLAG_NULLABLE

        return self

    def set_dictionary_ordered(self, dictionary_ordered) -> CSchemaBuilder:
        if dictionary_ordered:
            self._ptr.flags = self._ptr.flags | ARROW_FLAG_DICTIONARY_ORDERED
        else:
            self._ptr.flags = self._ptr.flags & ~ARROW_FLAG_DICTIONARY_ORDERED

        return self

    def set_map_keys_sorted(self, map_keys_sorted) -> CSchemaBuilder:
        if map_keys_sorted:
            self._ptr.flags = self._ptr.flags | ARROW_FLAG_MAP_KEYS_SORTED
        else:
            self._ptr.flags = self._ptr.flags & ~ARROW_FLAG_MAP_KEYS_SORTED

        return self

    def validate(self) -> CSchemaView:
        return CSchemaView(self.c_schema)

    def finish(self) -> CSchema:
        self.c_schema._assert_valid()
        cdef CSchema out = CSchema.allocate()
        ArrowSchemaMove(self.c_schema._ptr, out._ptr)
        ArrowSchemaInit(self.c_schema._ptr)
        return out
