# 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.

import struct
from typing import List, Optional, Any, Iterator, BinaryIO

import pyarrow as pa
import pyarrow.dataset as ds
from pyarrow import RecordBatch

from pypaimon.common.delta_varint_compressor import DeltaVarintCompressor
from pypaimon.common.file_io import FileIO
from pypaimon.common.map_blob_key_serializer import create_map_blob_key_serializer
from pypaimon.read.reader.iface.record_batch_reader import RecordBatchReader
from pypaimon.schema.data_types import (
    DataField,
    PyarrowFieldParser,
    AtomicType,
    MapType,
    is_array_blob_type,
    is_map_blob_type,
)
from pypaimon.table.row.blob import Blob
from pypaimon.table.row.generic_row import GenericRow
from pypaimon.table.row.row_kind import RowKind


class FormatBlobReader(RecordBatchReader):
    NULL_LENGTH = -1
    PLACE_HOLDER_LENGTH = -2

    def __init__(self, file_io: FileIO, file_path: str, read_fields: List[str],
                 full_fields: List[DataField], push_down_predicate: Any, blob_as_descriptor: bool,
                 batch_size: int = 1024, row_indices: Optional[Any] = None,
                 blob_parallelism: int = 1):
        self._file_io = file_io
        self._file_path = file_path
        self._push_down_predicate = push_down_predicate
        self._blob_as_descriptor = blob_as_descriptor
        self._batch_size = batch_size
        self._blob_parallelism = blob_parallelism

        # Initialize the low-level blob format reader
        self.file_path = file_path
        self.blob_lengths: List[int] = []
        self.blob_offsets: List[int] = []
        self.returned = False
        self._input_stream = None
        self._blob_iterator = None
        self._current_batch = None
        try:
            self._file_size = file_io.get_file_size(file_path)
            self._input_stream = file_io.new_input_stream(file_path)
            self._read_index()
            self._apply_row_indices(row_indices)

            # Set up fields and schema before deciding whether the stream can be dropped.
            if len(read_fields) > 1:
                raise RuntimeError("Blob reader only supports one field.")
            self._fields = read_fields
            full_fields_map = {field.name: field for field in full_fields}
            projected_data_fields = [full_fields_map[name] for name in read_fields]
            self._data_field = projected_data_fields[0]
            if isinstance(self._data_field.type, MapType) and not is_map_blob_type(
                self._data_field.type
            ):
                raise ValueError(
                    "Map-Blob value type must be BLOB, but is "
                    f"{self._data_field.type.value}."
                )
            self._is_array_blob = is_array_blob_type(self._data_field.type)
            self._is_map_blob = is_map_blob_type(self._data_field.type)
            self._schema = PyarrowFieldParser.from_paimon_schema(projected_data_fields)

            # Drop the shared stream: descriptor/concurrent reads yield BlobRefs
            # that each open their own stream (one stream isn't thread-safe).
            # Nested Blob formats need the stream to read their keys and indexes.
            if (
                not self._is_array_blob
                and not self._is_map_blob
                and (self._blob_as_descriptor or self._blob_parallelism > 1)
            ):
                self._input_stream.close()
                self._input_stream = None
        except Exception:
            self.close()
            raise

    def read_arrow_batch(self, start_idx=None, end_idx=None) -> Optional[RecordBatch]:
        """
         start_idx: start index record of the blob file
         end_idx: end index record of the blob file
        """
        if self._blob_iterator is None:
            if self.returned:
                return None
            self.returned = True
            batch_iterator = BlobRecordIterator(
                self._file_io, self.file_path, self.blob_lengths,
                self.blob_offsets, self._data_field, self._input_stream,
                blob_as_descriptor=(
                    self._blob_as_descriptor or self._blob_parallelism > 1
                )
            )
            self._blob_iterator = iter(batch_iterator)
        read_size = self._batch_size
        if start_idx is not None and end_idx is not None:
            if self._blob_iterator.current_position >= end_idx:
                return None
            if self._blob_iterator.current_position < start_idx:
                self._blob_iterator.current_position = start_idx
            read_size = min(end_idx - self._blob_iterator.current_position, self._batch_size)
        # Collect records for this batch
        pydict_data = {name: [] for name in self._fields}
        records_in_batch = 0
        blobs_to_resolve = []

        try:
            while True:
                blob_row = next(self._blob_iterator)
                if blob_row is None:
                    break
                blob = blob_row.values[0]
                for field_name in self._fields:
                    if self._is_map_blob:
                        row_index = len(pydict_data[field_name])
                        pydict_data[field_name].append(
                            self._map_value_for_arrow(
                                blob,
                                field_name,
                                row_index,
                                blobs_to_resolve,
                            )
                        )
                    elif self._is_array_blob:
                        row_index = len(pydict_data[field_name])
                        pydict_data[field_name].append(
                            self._array_value_for_arrow(
                                blob,
                                field_name,
                                row_index,
                                blobs_to_resolve,
                            )
                        )
                    else:
                        if blob is None:
                            pydict_data[field_name].append(None)
                        elif blob is Blob.PLACE_HOLDER:
                            raise RuntimeError(
                                "Blob placeholder is not supported by FormatBlobReader yet."
                            )
                        elif self._blob_as_descriptor:
                            pydict_data[field_name].append(blob.to_descriptor().serialize())
                        elif self._blob_parallelism > 1:
                            idx = len(pydict_data[field_name])
                            pydict_data[field_name].append(None)
                            blobs_to_resolve.append((field_name, idx, 'raw', None, blob))
                        else:
                            pydict_data[field_name].append(blob.to_data())

                records_in_batch += 1
                if records_in_batch >= read_size:
                    break

        except StopIteration:
            pass

        if blobs_to_resolve:
            self._resolve_blobs_concurrent(pydict_data, blobs_to_resolve)

        if records_in_batch == 0:
            return None

        # Create RecordBatch
        if self._push_down_predicate is None:
            # Convert to Table first, then to RecordBatch
            table = pa.Table.from_pydict(pydict_data, self._schema)
            if table.num_rows > 0:
                return table.to_batches()[0]
            else:
                return None
        else:
            # Apply predicate filtering
            pa_batch = pa.Table.from_pydict(pydict_data, self._schema)
            dataset = ds.InMemoryDataset(pa_batch)
            scanner = dataset.scanner(filter=self._push_down_predicate)
            combine_chunks = scanner.to_table().combine_chunks()
            if combine_chunks.num_rows > 0:
                return combine_chunks.to_batches()[0]
            else:
                return None

    def _resolve_blobs_concurrent(self, pydict_data, blobs_to_resolve):
        blobs = [item[4] for item in blobs_to_resolve]
        results = self._file_io.read_blobs_concurrent(blobs, self._blob_parallelism)
        for target, data in zip(blobs_to_resolve, results):
            field_name, row_index, container_kind, slot, _ = target
            if container_kind == 'raw':
                pydict_data[field_name][row_index] = data
            elif container_kind == 'array':
                pydict_data[field_name][row_index][slot] = data
            elif container_kind == 'map':
                key = pydict_data[field_name][row_index][slot][0]
                pydict_data[field_name][row_index][slot] = (key, data)

    def _array_value_for_arrow(
        self,
        blob_array,
        field_name,
        row_index,
        blobs_to_resolve,
    ):
        if blob_array is None:
            return None
        if blob_array is Blob.ARRAY_PLACE_HOLDER:
            raise RuntimeError(
                "Blob placeholder is not supported by FormatBlobReader yet."
            )
        result = []
        for element_index, blob in enumerate(blob_array):
            if blob is None:
                result.append(None)
            elif self._blob_as_descriptor:
                result.append(blob.to_descriptor().serialize())
            elif self._blob_parallelism > 1:
                result.append(None)
                blobs_to_resolve.append(
                    (field_name, row_index, 'array', element_index, blob)
                )
            else:
                result.append(blob.to_data())
        return result

    def _map_value_for_arrow(
        self,
        blob_map,
        field_name,
        row_index,
        blobs_to_resolve,
    ):
        if blob_map is None:
            return None
        if blob_map is Blob.MAP_PLACE_HOLDER:
            raise RuntimeError(
                "Blob placeholder is not supported by FormatBlobReader yet."
            )
        if None in blob_map:
            raise ValueError(
                "MAP<X, BLOB> with null keys cannot be converted to a PyArrow Map."
            )

        result = []
        for entry_index, (key, blob) in enumerate(blob_map.items()):
            if blob is None:
                value = None
            elif self._blob_as_descriptor:
                value = blob.to_descriptor().serialize()
            elif self._blob_parallelism > 1:
                value = None
                blobs_to_resolve.append(
                    (field_name, row_index, 'map', entry_index, blob)
                )
            else:
                value = blob.to_data()
            result.append((key, value))
        return result

    def close(self):
        self._blob_iterator = None
        if self._input_stream is not None:
            self._input_stream.close()
            self._input_stream = None

    def _read_index(self) -> None:
        f = self._input_stream

        # Seek to header: last 5 bytes
        f.seek(self._file_size - 5)
        header = f.read(5)

        if len(header) != 5:
            raise IOError("Invalid blob file: cannot read header")

        # Parse header
        index_length = struct.unpack('<I', header[:4])[0]  # Little endian
        version = header[4]

        if version != 1:
            raise IOError(f"Unsupported blob file version: {version}")

        # Read index data
        f.seek(self._file_size - 5 - index_length)
        index_bytes = f.read(index_length)

        if len(index_bytes) != index_length:
            raise IOError("Invalid blob file: cannot read index")

        # Decompress blob lengths and compute offsets
        blob_lengths = DeltaVarintCompressor.decompress(index_bytes)
        blob_offsets = []
        offset = 0
        for length in blob_lengths:
            if length < 0:
                blob_offsets.append(-1)
            else:
                blob_offsets.append(offset)
                offset += length
        self.blob_lengths = blob_lengths
        self.blob_offsets = blob_offsets

    def _apply_row_indices(self, row_indices: Optional[Any]) -> None:
        if row_indices is None:
            return

        selected_lengths = []
        selected_offsets = []
        record_count = len(self.blob_lengths)
        for row_index in row_indices:
            row_index = int(row_index)
            if row_index < 0 or row_index >= record_count:
                raise IndexError(
                    f"Blob row index {row_index} is out of range for file "
                    f"{self.file_path}, record count: {record_count}."
                )
            selected_lengths.append(self.blob_lengths[row_index])
            selected_offsets.append(self.blob_offsets[row_index])

        self.blob_lengths = selected_lengths
        self.blob_offsets = selected_offsets


class BlobRecordIterator:
    MAGIC_NUMBER_SIZE = 4
    METADATA_OVERHEAD = 16
    ARRAY_HEADER_SIZE = 9
    ARRAY_VERSION = 1
    ARRAY_MAGIC_NUMBER = 1094861634
    ARRAY_NULL_ELEMENT_LENGTH = -1
    ARRAY_INDEX_LENGTH_SIZE = 4
    MIN_ARRAY_PAYLOAD_LENGTH = ARRAY_HEADER_SIZE + ARRAY_INDEX_LENGTH_SIZE
    MAP_HEADER_SIZE = 9
    MAP_VERSION = 1
    MAP_MAGIC_NUMBER = 0x4D424342
    MAP_NULL_KEY_LENGTH = -1
    MAP_NULL_VALUE_LENGTH = -1
    MAP_INDEX_LENGTH_SIZE = 4
    MAP_INDEX_LENGTHS_SIZE = MAP_INDEX_LENGTH_SIZE * 2
    MIN_MAP_PAYLOAD_LENGTH = MAP_HEADER_SIZE + MAP_INDEX_LENGTHS_SIZE
    NULL_LENGTH = -1
    PLACE_HOLDER_LENGTH = -2

    def __init__(self, file_io: FileIO, file_path: str, blob_lengths: List[int],
                 blob_offsets: List[int], field,
                 input_stream: Optional[BinaryIO] = None,
                 blob_as_descriptor: bool = False):
        self.file_io = file_io
        self.file_path = file_path
        self.input_stream = input_stream
        if isinstance(field, DataField):
            self.field = field
        else:
            self.field = DataField(0, field, AtomicType("BLOB"))
        self.field_name = self.field.name
        if isinstance(self.field.type, MapType) and not is_map_blob_type(self.field.type):
            raise ValueError(
                f"Map-Blob value type must be BLOB, but is {self.field.type.value}."
            )
        self.is_array_blob = is_array_blob_type(self.field.type)
        self.is_map_blob = is_map_blob_type(self.field.type)
        self.map_key_serializer = (
            create_map_blob_key_serializer(self.field.type.key)
            if self.is_map_blob
            else None
        )
        self.blob_as_descriptor = blob_as_descriptor
        self.blob_lengths = blob_lengths
        self.blob_offsets = blob_offsets
        self.current_position = 0

    def __iter__(self) -> Iterator[GenericRow]:
        return self

    def __next__(self) -> GenericRow:
        if self.current_position >= len(self.blob_lengths):
            raise StopIteration
        fields = [self.field]
        length = self.blob_lengths[self.current_position]
        if length == self.NULL_LENGTH:
            self.current_position += 1
            return GenericRow([None], fields, RowKind.INSERT)
        if length == self.PLACE_HOLDER_LENGTH:
            self.current_position += 1
            if self.is_map_blob:
                placeholder = Blob.MAP_PLACE_HOLDER
            elif self.is_array_blob:
                placeholder = Blob.ARRAY_PLACE_HOLDER
            else:
                placeholder = Blob.PLACE_HOLDER
            return GenericRow([placeholder], fields, RowKind.INSERT)
        # Create blob reference for the current blob
        # Skip magic number (4 bytes) and exclude length (8 bytes) + CRC (4 bytes) = 12 bytes
        blob_offset = self.blob_offsets[self.current_position] + self.MAGIC_NUMBER_SIZE  # Skip magic number
        blob_length = length - self.METADATA_OVERHEAD
        if self.is_map_blob:
            blob = self._read_blob_map(blob_offset, blob_length)
        elif self.is_array_blob:
            blob = self._read_blob_array(blob_offset, blob_length)
        elif self.input_stream is not None and not self.blob_as_descriptor:
            blob = Blob.from_data(self._read_inline_blob(blob_offset, blob_length))
        else:
            blob = Blob.from_file(self.file_io, self.file_path, blob_offset, blob_length)
        self.current_position += 1
        return GenericRow([blob], fields, RowKind.INSERT)

    def returned_position(self) -> int:
        return self.current_position

    def _read_inline_blob(self, position: int, length: int) -> bytes:
        self.input_stream.seek(position)
        data = self._read_fully(length)
        if len(data) != length:
            raise IOError("Invalid blob file: cannot read blob data")
        return data

    def _read_blob_array(self, position: int, length: int):
        if position < 0 or length < self.MIN_ARRAY_PAYLOAD_LENGTH:
            raise ValueError(
                f"Invalid ARRAY<BLOB> payload position or length: {position}, {length}"
            )

        stream = self.input_stream
        close_stream = False
        if stream is None:
            stream = self.file_io.new_input_stream(self.file_path)
            close_stream = True
        try:
            stream.seek(position)
            header = self._read_fully_from(stream, self.ARRAY_HEADER_SIZE)
            if len(header) != self.ARRAY_HEADER_SIZE:
                raise IOError("Invalid ARRAY<BLOB> payload: cannot read header")
            magic, version, element_count = struct.unpack('<IBI', header)
            if magic != self.ARRAY_MAGIC_NUMBER:
                raise ValueError(f"Invalid ARRAY<BLOB> payload magic number: {magic}")
            if version != self.ARRAY_VERSION:
                raise ValueError(f"Unsupported ARRAY<BLOB> payload version: {version}")
            if element_count > 0x7fffffff:
                raise ValueError(f"Invalid ARRAY<BLOB> element count: {element_count}")

            payload_end = position + length
            element_data_start = position + self.ARRAY_HEADER_SIZE
            index_length_position = payload_end - self.ARRAY_INDEX_LENGTH_SIZE
            stream.seek(index_length_position)
            index_length_bytes = self._read_fully_from(stream, self.ARRAY_INDEX_LENGTH_SIZE)
            if len(index_length_bytes) != self.ARRAY_INDEX_LENGTH_SIZE:
                raise IOError("Invalid ARRAY<BLOB> payload: cannot read index length")
            index_length = struct.unpack('<I', index_length_bytes)[0]
            maximum_index_length = length - self.MIN_ARRAY_PAYLOAD_LENGTH
            if index_length > 0x7fffffff or index_length > maximum_index_length:
                raise ValueError(
                    f"Invalid ARRAY<BLOB> element index length: {index_length}"
                )
            if element_count > index_length:
                raise ValueError(
                    "ARRAY<BLOB> element count exceeds element index length."
                )

            element_index_start = index_length_position - index_length
            stream.seek(element_index_start)
            index_bytes = self._read_fully_from(stream, index_length)
            if len(index_bytes) != index_length:
                raise IOError("Invalid ARRAY<BLOB> payload: cannot read element index")
            self._validate_array_element_index(index_bytes)
            element_lengths = DeltaVarintCompressor.decompress(index_bytes)
            if len(element_lengths) != element_count:
                raise ValueError(
                    "ARRAY<BLOB> element count does not match element index length."
                )

            element_data_length = element_index_start - element_data_start
            total_element_length = 0
            for element_length in element_lengths:
                if element_length == self.ARRAY_NULL_ELEMENT_LENGTH:
                    continue
                if element_length < 0:
                    raise ValueError(
                        f"Invalid ARRAY<BLOB> element length: {element_length}"
                    )
                if element_length > element_data_length - total_element_length:
                    raise ValueError(
                        "ARRAY<BLOB> element lengths exceed the payload data length."
                    )
                total_element_length += element_length
            if total_element_length != element_data_length:
                raise ValueError(
                    "ARRAY<BLOB> element lengths do not match the payload data length."
                )

            element_data = None
            if not self.blob_as_descriptor:
                stream.seek(element_data_start)
                element_data = self._read_fully_from(stream, element_data_length)
                if len(element_data) != element_data_length:
                    raise IOError("Invalid ARRAY<BLOB> payload: cannot read element data")

            blobs = []
            element_offset = element_data_start
            data_offset = 0
            for element_length in element_lengths:
                if element_length == self.ARRAY_NULL_ELEMENT_LENGTH:
                    blobs.append(None)
                    continue
                if self.blob_as_descriptor:
                    blobs.append(
                        Blob.from_file(self.file_io, self.file_path, element_offset, element_length)
                    )
                else:
                    blobs.append(Blob.from_data(
                        element_data[data_offset:data_offset + element_length]
                    ))
                element_offset += element_length
                data_offset += element_length
            return blobs
        finally:
            if close_stream:
                stream.close()

    def _read_blob_map(self, position: int, length: int):
        if position < 0 or length < self.MIN_MAP_PAYLOAD_LENGTH:
            raise ValueError(
                f"Invalid MAP<X, BLOB> payload position or length: {position}, {length}"
            )

        stream = self.input_stream
        close_stream = False
        if stream is None:
            stream = self.file_io.new_input_stream(self.file_path)
            close_stream = True
        try:
            stream.seek(position)
            header = self._read_fully_from(stream, self.MAP_HEADER_SIZE)
            if len(header) != self.MAP_HEADER_SIZE:
                raise IOError("Invalid MAP<X, BLOB> payload: cannot read header")
            magic, version, entry_count = struct.unpack('<IBI', header)
            if magic != self.MAP_MAGIC_NUMBER:
                raise ValueError(f"Invalid MAP<X, BLOB> payload magic number: {magic}")
            if version != self.MAP_VERSION:
                raise ValueError(f"Unsupported MAP<X, BLOB> payload version: {version}")
            if entry_count > 0x7fffffff:
                raise ValueError(f"Invalid MAP<X, BLOB> entry count: {entry_count}")

            payload_end = position + length
            data_start = position + self.MAP_HEADER_SIZE
            index_lengths_position = payload_end - self.MAP_INDEX_LENGTHS_SIZE
            stream.seek(index_lengths_position)
            index_length_bytes = self._read_fully_from(
                stream, self.MAP_INDEX_LENGTHS_SIZE
            )
            if len(index_length_bytes) != self.MAP_INDEX_LENGTHS_SIZE:
                raise IOError("Invalid MAP<X, BLOB> payload: cannot read index lengths")
            key_index_length, value_index_length = struct.unpack('<II', index_length_bytes)
            self._check_map_index_lengths(
                key_index_length,
                value_index_length,
                length,
                entry_count,
            )

            value_index_start = index_lengths_position - value_index_length
            key_index_start = value_index_start - key_index_length
            stream.seek(key_index_start)
            key_index_bytes = self._read_fully_from(stream, key_index_length)
            if len(key_index_bytes) != key_index_length:
                raise IOError("Invalid MAP<X, BLOB> payload: cannot read key index")
            stream.seek(value_index_start)
            value_index_bytes = self._read_fully_from(stream, value_index_length)
            if len(value_index_bytes) != value_index_length:
                raise IOError("Invalid MAP<X, BLOB> payload: cannot read value index")

            self._validate_map_index(key_index_bytes, "key")
            self._validate_map_index(value_index_bytes, "value")
            try:
                key_lengths = DeltaVarintCompressor.decompress(key_index_bytes)
            except RuntimeError as error:
                raise ValueError("Invalid MAP<X, BLOB> key index.") from error
            try:
                value_lengths = DeltaVarintCompressor.decompress(value_index_bytes)
            except RuntimeError as error:
                raise ValueError("Invalid MAP<X, BLOB> value index.") from error

            data_length = key_index_start - data_start
            key_data_length = self._check_map_key_lengths(
                key_lengths, data_length, entry_count
            )
            value_data_length = data_length - key_data_length
            self._check_map_value_lengths(
                value_lengths, value_data_length, entry_count
            )

            stream.seek(data_start)
            key_data = self._read_fully_from(stream, key_data_length)
            if len(key_data) != key_data_length:
                raise IOError("Invalid MAP<X, BLOB> payload: cannot read key data")
            keys = []
            key_data_offset = 0
            for key_length in key_lengths:
                if key_length == self.MAP_NULL_KEY_LENGTH:
                    keys.append(None)
                    continue
                serialized_key = key_data[
                    key_data_offset:key_data_offset + key_length
                ]
                try:
                    keys.append(self.map_key_serializer.deserialize(serialized_key))
                except ValueError as error:
                    raise ValueError("Invalid MAP<X, BLOB> key.") from error
                key_data_offset += key_length

            value_data_start = data_start + key_data_length
            value_data = None
            if not self.blob_as_descriptor:
                stream.seek(value_data_start)
                value_data = self._read_fully_from(stream, value_data_length)
                if len(value_data) != value_data_length:
                    raise IOError("Invalid MAP<X, BLOB> payload: cannot read value data")

            result = {}
            value_offset = value_data_start
            value_data_offset = 0
            for key, value_length in zip(keys, value_lengths):
                if value_length == self.MAP_NULL_VALUE_LENGTH:
                    value = None
                elif self.blob_as_descriptor:
                    value = Blob.from_file(
                        self.file_io, self.file_path, value_offset, value_length
                    )
                else:
                    value = Blob.from_data(
                        value_data[value_data_offset:value_data_offset + value_length]
                    )
                if value_length != self.MAP_NULL_VALUE_LENGTH:
                    value_offset += value_length
                    value_data_offset += value_length
                result[key] = value
            return result
        finally:
            if close_stream:
                stream.close()

    @classmethod
    def _check_map_index_lengths(
        cls,
        key_index_length: int,
        value_index_length: int,
        payload_length: int,
        entry_count: int,
    ) -> None:
        maximum_indexes_length = payload_length - cls.MIN_MAP_PAYLOAD_LENGTH
        if key_index_length > 0x7fffffff or key_index_length > maximum_indexes_length:
            raise ValueError(
                f"Invalid MAP<X, BLOB> key index length: {key_index_length}"
            )
        if value_index_length > 0x7fffffff or value_index_length > maximum_indexes_length:
            raise ValueError(
                f"Invalid MAP<X, BLOB> value index length: {value_index_length}"
            )
        if key_index_length + value_index_length > maximum_indexes_length:
            raise ValueError("MAP<X, BLOB> indexes exceed the payload length.")
        if entry_count > key_index_length:
            raise ValueError("MAP<X, BLOB> entry count exceeds key index length.")
        if entry_count > value_index_length:
            raise ValueError("MAP<X, BLOB> entry count exceeds value index length.")

    def _check_map_key_lengths(
        self, key_lengths, max_data_length: int, entry_count: int
    ) -> int:
        if len(key_lengths) != entry_count:
            raise ValueError(
                "MAP<X, BLOB> entry count does not match key index length."
            )

        key_data_length = 0
        fixed_key_length = self.map_key_serializer.fixed_length
        for key_length in key_lengths:
            if key_length == self.MAP_NULL_KEY_LENGTH:
                continue
            if key_length < 0:
                raise ValueError(f"Invalid MAP<X, BLOB> key length: {key_length}")
            if key_length > 0x7fffffff:
                raise ValueError(f"MAP<X, BLOB> key is too large: {key_length}")
            if fixed_key_length >= 0 and key_length != fixed_key_length:
                raise ValueError(
                    f"Invalid MAP<X, BLOB> fixed-width key length: {key_length}"
                )
            if key_length > max_data_length - key_data_length:
                raise ValueError(
                    "MAP<X, BLOB> key lengths exceed the payload data length."
                )
            key_data_length += key_length
        return key_data_length

    def _check_map_value_lengths(
        self, value_lengths, max_data_length: int, entry_count: int
    ) -> None:
        if len(value_lengths) != entry_count:
            raise ValueError(
                "MAP<X, BLOB> entry count does not match value index length."
            )

        value_data_length = 0
        for value_length in value_lengths:
            if value_length == self.MAP_NULL_VALUE_LENGTH:
                continue
            if value_length < 0:
                raise ValueError(f"Invalid MAP<X, BLOB> value length: {value_length}")
            if not self.blob_as_descriptor and value_length > 0x7fffffff:
                raise ValueError(
                    f"MAP<X, BLOB> inline value is too large: {value_length}"
                )
            if value_length > max_data_length - value_data_length:
                raise ValueError(
                    "MAP<X, BLOB> value lengths exceed the payload data length."
                )
            value_data_length += value_length
        if value_data_length != max_data_length:
            raise ValueError(
                "MAP<X, BLOB> key/value lengths do not match the payload data length."
            )

    @staticmethod
    def _validate_map_index(index_bytes: bytes, index_name: str) -> None:
        varint_length = 0
        for value in index_bytes:
            varint_length += 1
            if varint_length > 10:
                raise ValueError(f"Invalid MAP<X, BLOB> {index_name} index.")
            if value & 0x80 == 0:
                varint_length = 0
        if varint_length != 0:
            raise ValueError(f"Invalid MAP<X, BLOB> {index_name} index.")

    @staticmethod
    def _validate_array_element_index(index_bytes: bytes) -> None:
        varint_length = 0
        for value in index_bytes:
            varint_length += 1
            if varint_length > 10:
                raise ValueError("Invalid ARRAY<BLOB> element index.")
            if value & 0x80 == 0:
                varint_length = 0
        if varint_length != 0:
            raise ValueError("Invalid ARRAY<BLOB> element index.")

    def _read_fully(self, length: int) -> bytes:
        return self._read_fully_from(self.input_stream, length)

    @staticmethod
    def _read_fully_from(stream, length: int) -> bytes:
        data = bytearray()
        while len(data) < length:
            chunk = stream.read(length - len(data))
            if not chunk:
                break
            data.extend(chunk)
        return bytes(data)
