#
# 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 base64
import decimal
import datetime
import json
import struct
from array import array
from typing import Any, Callable, Dict, List, NamedTuple, Tuple
from pyspark.errors import PySparkValueError
from zoneinfo import ZoneInfo


class VariantUtils:
    """
    A utility class for VariantVal.

    Adapted from library at: org.apache.spark.types.variant.VariantUtil
    """

    BASIC_TYPE_BITS = 2
    BASIC_TYPE_MASK = 0x3
    TYPE_INFO_MASK = 0x3F
    # The inclusive maximum value of the type info value. It is the size limit of `SHORT_STR`.
    MAX_SHORT_STR_SIZE = 0x3F

    # Below is all possible basic type values.
    # Primitive value. The type info value must be one of the values in the below section.
    PRIMITIVE = 0
    # Short string value. The type info value is the string size, which must be in `[0,
    # MAX_SHORT_STR_SIZE]`.
    # The string content bytes directly follow the header byte.
    SHORT_STR = 1
    # Object value. The content contains a size, a list of field ids, a list of field offsets, and
    # the actual field data. The length of the id list is `size`, while the length of the offset
    # list is `size + 1`, where the last offset represent the total size of the field data. The
    # fields in an object must be sorted by the field name in alphabetical order. Duplicate field
    # names in one object are not allowed.
    # We use 5 bits in the type info to specify the integer type of the object header: it should
    # be 0_b4_b3b2_b1b0 (MSB is 0), where:
    # - b4 specifies the type of size. When it is 0/1, `size` is a little-endian 1/4-byte
    # unsigned integer.
    # - b3b2/b1b0 specifies the integer type of id and offset. When the 2 bits are  0/1/2, the
    # list contains 1/2/3-byte little-endian unsigned integers.
    OBJECT = 2
    # Array value. The content contains a size, a list of field offsets, and the actual element
    # data. It is similar to an object without the id list. The length of the offset list
    # is `size + 1`, where the last offset represent the total size of the element data.
    # Its type info should be: 000_b2_b1b0:
    # - b2 specifies the type of size.
    # - b1b0 specifies the integer type of offset.
    ARRAY = 3

    # Below is all possible type info values for `PRIMITIVE`.
    # JSON Null value. Empty content.
    NULL = 0
    # True value. Empty content.
    TRUE = 1
    # False value. Empty content.
    FALSE = 2
    # 1-byte little-endian signed integer.
    INT1 = 3
    # 2-byte little-endian signed integer.
    INT2 = 4
    # 4-byte little-endian signed integer.
    INT4 = 5
    # 4-byte little-endian signed integer.
    INT8 = 6
    # 8-byte IEEE double.
    DOUBLE = 7
    # 4-byte decimal. Content is 1-byte scale + 4-byte little-endian signed integer.
    DECIMAL4 = 8
    # 8-byte decimal. Content is 1-byte scale + 8-byte little-endian signed integer.
    DECIMAL8 = 9
    # 16-byte decimal. Content is 1-byte scale + 16-byte little-endian signed integer.
    DECIMAL16 = 10
    # Date value. Content is 4-byte little-endian signed integer that represents the number of days
    # from the Unix epoch.
    DATE = 11
    # Timestamp value. Content is 8-byte little-endian signed integer that represents the number of
    # microseconds elapsed since the Unix epoch, 1970-01-01 00:00:00 UTC. This is a timezone-aware
    # field and when reading into a Python datetime object defaults to the UTC timezone.
    TIMESTAMP = 12
    # Timestamp_ntz value. It has the same content as `TIMESTAMP` but should always be interpreted
    # as if the local time zone is UTC.
    TIMESTAMP_NTZ = 13
    # 4-byte IEEE float.
    FLOAT = 14
    # Binary value. The content is (4-byte little-endian unsigned integer representing the binary
    # size) + (size bytes of binary content).
    BINARY = 15
    # Long string value. The content is (4-byte little-endian unsigned integer representing the
    # string size) + (size bytes of string content).
    LONG_STR = 16

    VERSION = 1
    # The lower 4 bits of the first metadata byte contain the version.
    VERSION_MASK = 0x0F

    U8_MAX = 0xFF
    U16_MAX = 0xFFFF
    U24_MAX = 0xFFFFFF
    U24_SIZE = 3
    U32_SIZE = 4

    I8_MAX = 0x7F
    I8_MIN = -0x80
    I16_MAX = 0x7FFF
    I16_MIN = -0x8000
    I32_MAX = 0x7FFFFFFF
    I32_MIN = -0x80000000
    I64_MAX = 0x7FFFFFFFFFFFFFFF
    I64_MIN = -0x8000000000000000

    EPOCH = datetime.datetime(
        year=1970, month=1, day=1, hour=0, minute=0, second=0, tzinfo=datetime.timezone.utc
    )
    EPOCH_NTZ = datetime.datetime(year=1970, month=1, day=1, hour=0, minute=0, second=0)

    MAX_DECIMAL4_PRECISION = 9
    MAX_DECIMAL4_VALUE = 10**MAX_DECIMAL4_PRECISION
    MAX_DECIMAL8_PRECISION = 18
    MAX_DECIMAL8_VALUE = 10**MAX_DECIMAL8_PRECISION
    MAX_DECIMAL16_PRECISION = 38
    MAX_DECIMAL16_VALUE = 10**MAX_DECIMAL16_PRECISION

    @classmethod
    def to_json(cls, value: bytes, metadata: bytes, zone_id: str = "UTC") -> str:
        """
        Convert the VariantVal to a JSON string. The `zone_id` parameter denotes the time zone that
        timestamp fields should be parsed in. It defaults to "UTC". The list of valid zone IDs can
        found by importing the `zoneinfo` module and running `zoneinfo.available_timezones()`.
        :return: JSON string
        """
        return cls._to_json(value, metadata, 0, zone_id)

    @classmethod
    def to_python(cls, value: bytes, metadata: bytes) -> str:
        """
        Convert the VariantVal to a nested Python object of Python data types.
        :return: Python representation of the Variant nested structure
        """
        return cls._to_python(value, metadata, 0)

    @classmethod
    def parse_json(cls, json_str: str) -> Tuple[bytes, bytes]:
        """
        Parses the JSON string and creates the Variant binary (value, metadata)
        :return: tuple of 2 binary values (value, metadata)
        """
        builder = VariantBuilder()
        return builder.build(json_str)

    @classmethod
    def _read_long(cls, data: bytes, pos: int, num_bytes: int, signed: bool) -> int:
        cls._check_index(pos, len(data))
        cls._check_index(pos + num_bytes - 1, len(data))
        return int.from_bytes(data[pos : pos + num_bytes], byteorder="little", signed=signed)

    @classmethod
    def _check_index(cls, pos: int, length: int) -> None:
        if pos < 0 or pos >= length:
            raise PySparkValueError(errorClass="MALFORMED_VARIANT", messageParameters={})

    @classmethod
    def _get_type_info(cls, value: bytes, pos: int) -> Tuple[int, int]:
        """
        Returns the (basic_type, type_info) pair from the given position in the value.
        """
        basic_type = value[pos] & VariantUtils.BASIC_TYPE_MASK
        type_info = (value[pos] >> VariantUtils.BASIC_TYPE_BITS) & VariantUtils.TYPE_INFO_MASK
        return (basic_type, type_info)

    @classmethod
    def _get_metadata_key(cls, metadata: bytes, id: int) -> str:
        """
        Returns the key string from the dictionary in the metadata, corresponding to `id`.
        """
        cls._check_index(0, len(metadata))
        offset_size = ((metadata[0] >> 6) & 0x3) + 1
        dict_size = cls._read_long(metadata, 1, offset_size, signed=False)
        if id >= dict_size:
            raise PySparkValueError(errorClass="MALFORMED_VARIANT", messageParameters={})
        string_start = 1 + (dict_size + 2) * offset_size
        offset = cls._read_long(metadata, 1 + (id + 1) * offset_size, offset_size, signed=False)
        next_offset = cls._read_long(
            metadata, 1 + (id + 2) * offset_size, offset_size, signed=False
        )
        if offset > next_offset:
            raise PySparkValueError(errorClass="MALFORMED_VARIANT", messageParameters={})
        cls._check_index(string_start + next_offset - 1, len(metadata))
        return metadata[string_start + offset : (string_start + next_offset)].decode("utf-8")

    @classmethod
    def _get_boolean(cls, value: bytes, pos: int) -> bool:
        cls._check_index(pos, len(value))
        basic_type, type_info = cls._get_type_info(value, pos)
        if basic_type != VariantUtils.PRIMITIVE or (
            type_info != VariantUtils.TRUE and type_info != VariantUtils.FALSE
        ):
            raise PySparkValueError(errorClass="MALFORMED_VARIANT", messageParameters={})
        return type_info == VariantUtils.TRUE

    @classmethod
    def _get_long(cls, value: bytes, pos: int) -> int:
        cls._check_index(pos, len(value))
        basic_type, type_info = cls._get_type_info(value, pos)
        if basic_type != VariantUtils.PRIMITIVE:
            raise PySparkValueError(errorClass="MALFORMED_VARIANT", messageParameters={})
        if type_info == VariantUtils.INT1:
            return cls._read_long(value, pos + 1, 1, signed=True)
        elif type_info == VariantUtils.INT2:
            return cls._read_long(value, pos + 1, 2, signed=True)
        elif type_info == VariantUtils.INT4 or type_info == VariantUtils.DATE:
            return cls._read_long(value, pos + 1, 4, signed=True)
        elif type_info == VariantUtils.INT8:
            return cls._read_long(value, pos + 1, 8, signed=True)
        raise PySparkValueError(errorClass="MALFORMED_VARIANT", messageParameters={})

    @classmethod
    def _get_date(cls, value: bytes, pos: int) -> datetime.date:
        cls._check_index(pos, len(value))
        basic_type, type_info = cls._get_type_info(value, pos)
        if basic_type != VariantUtils.PRIMITIVE:
            raise PySparkValueError(errorClass="MALFORMED_VARIANT", messageParameters={})
        if type_info == VariantUtils.DATE:
            days_since_epoch = cls._read_long(value, pos + 1, 4, signed=True)
            return datetime.date.fromordinal(VariantUtils.EPOCH.toordinal() + days_since_epoch)
        raise PySparkValueError(errorClass="MALFORMED_VARIANT", messageParameters={})

    @classmethod
    def _get_timestamp(cls, value: bytes, pos: int, zone_id: str) -> datetime.datetime:
        cls._check_index(pos, len(value))
        basic_type, type_info = cls._get_type_info(value, pos)
        if basic_type != VariantUtils.PRIMITIVE:
            raise PySparkValueError(errorClass="MALFORMED_VARIANT", messageParameters={})
        if type_info == VariantUtils.TIMESTAMP_NTZ:
            microseconds_since_epoch = cls._read_long(value, pos + 1, 8, signed=True)
            return VariantUtils.EPOCH_NTZ + datetime.timedelta(
                microseconds=microseconds_since_epoch
            )
        if type_info == VariantUtils.TIMESTAMP:
            microseconds_since_epoch = cls._read_long(value, pos + 1, 8, signed=True)
            return (
                VariantUtils.EPOCH + datetime.timedelta(microseconds=microseconds_since_epoch)
            ).astimezone(ZoneInfo(zone_id))
        raise PySparkValueError(errorClass="MALFORMED_VARIANT", messageParameters={})

    @classmethod
    def _get_string(cls, value: bytes, pos: int) -> str:
        cls._check_index(pos, len(value))
        basic_type, type_info = cls._get_type_info(value, pos)
        if basic_type == VariantUtils.SHORT_STR or (
            basic_type == VariantUtils.PRIMITIVE and type_info == VariantUtils.LONG_STR
        ):
            start = 0
            length = 0
            if basic_type == VariantUtils.SHORT_STR:
                start = pos + 1
                length = type_info
            else:
                start = pos + 1 + VariantUtils.U32_SIZE
                length = cls._read_long(value, pos + 1, VariantUtils.U32_SIZE, signed=False)
            cls._check_index(start + length - 1, len(value))
            return value[start : start + length].decode("utf-8")
        raise PySparkValueError(errorClass="MALFORMED_VARIANT", messageParameters={})

    @classmethod
    def _get_double(cls, value: bytes, pos: int) -> float:
        cls._check_index(pos, len(value))
        basic_type, type_info = cls._get_type_info(value, pos)
        if basic_type != VariantUtils.PRIMITIVE:
            raise PySparkValueError(errorClass="MALFORMED_VARIANT", messageParameters={})
        if type_info == VariantUtils.FLOAT:
            cls._check_index(pos + 4, len(value))
            return struct.unpack("<f", value[pos + 1 : pos + 5])[0]
        elif type_info == VariantUtils.DOUBLE:
            cls._check_index(pos + 8, len(value))
            return struct.unpack("<d", value[pos + 1 : pos + 9])[0]
        raise PySparkValueError(errorClass="MALFORMED_VARIANT", messageParameters={})

    @classmethod
    def _check_decimal(cls, unscaled: int, scale: int, max_unscaled: int, max_scale: int) -> None:
        # max_unscaled == 10**max_scale, but we pass a literal parameter to avoid redundant
        # computation.
        if unscaled >= max_unscaled or unscaled <= -max_unscaled or scale > max_scale:
            raise PySparkValueError(errorClass="MALFORMED_VARIANT", messageParameters={})

    @classmethod
    def _get_decimal(cls, value: bytes, pos: int) -> decimal.Decimal:
        cls._check_index(pos, len(value))
        basic_type, type_info = cls._get_type_info(value, pos)
        if basic_type != VariantUtils.PRIMITIVE:
            raise PySparkValueError(errorClass="MALFORMED_VARIANT", messageParameters={})
        scale = value[pos + 1]
        unscaled = 0
        if type_info == VariantUtils.DECIMAL4:
            unscaled = cls._read_long(value, pos + 2, 4, signed=True)
            cls._check_decimal(unscaled, scale, cls.MAX_DECIMAL4_VALUE, cls.MAX_DECIMAL4_PRECISION)
        elif type_info == VariantUtils.DECIMAL8:
            unscaled = cls._read_long(value, pos + 2, 8, signed=True)
            cls._check_decimal(unscaled, scale, cls.MAX_DECIMAL8_VALUE, cls.MAX_DECIMAL8_PRECISION)
        elif type_info == VariantUtils.DECIMAL16:
            cls._check_index(pos + 17, len(value))
            unscaled = int.from_bytes(value[pos + 2 : pos + 18], byteorder="little", signed=True)
            cls._check_decimal(
                unscaled, scale, cls.MAX_DECIMAL16_VALUE, cls.MAX_DECIMAL16_PRECISION
            )
        else:
            raise PySparkValueError(errorClass="MALFORMED_VARIANT", messageParameters={})
        return decimal.Decimal(unscaled) * (decimal.Decimal(10) ** (-scale))

    @classmethod
    def _get_binary(cls, value: bytes, pos: int) -> bytes:
        cls._check_index(pos, len(value))
        basic_type, type_info = cls._get_type_info(value, pos)
        if basic_type != VariantUtils.PRIMITIVE or type_info != VariantUtils.BINARY:
            raise PySparkValueError(errorClass="MALFORMED_VARIANT", messageParameters={})
        start = pos + 1 + VariantUtils.U32_SIZE
        length = cls._read_long(value, pos + 1, VariantUtils.U32_SIZE, signed=False)
        cls._check_index(start + length - 1, len(value))
        return bytes(value[start : start + length])

    @classmethod
    def _get_type(cls, value: bytes, pos: int) -> Any:
        """
        Returns the Python type of the Variant at the given position.
        """
        cls._check_index(pos, len(value))
        basic_type, type_info = cls._get_type_info(value, pos)
        if basic_type == VariantUtils.SHORT_STR:
            return str
        elif basic_type == VariantUtils.OBJECT:
            return dict
        elif basic_type == VariantUtils.ARRAY:
            return array
        elif type_info == VariantUtils.NULL:
            return type(None)
        elif type_info == VariantUtils.TRUE or type_info == VariantUtils.FALSE:
            return bool
        elif (
            type_info == VariantUtils.INT1
            or type_info == VariantUtils.INT2
            or type_info == VariantUtils.INT4
            or type_info == VariantUtils.INT8
        ):
            return int
        elif type_info == VariantUtils.DOUBLE or type_info == VariantUtils.FLOAT:
            return float
        elif (
            type_info == VariantUtils.DECIMAL4
            or type_info == VariantUtils.DECIMAL8
            or type_info == VariantUtils.DECIMAL16
        ):
            return decimal.Decimal
        elif type_info == VariantUtils.BINARY:
            return bytes
        elif type_info == VariantUtils.DATE:
            return datetime.date
        elif type_info == VariantUtils.TIMESTAMP or type_info == VariantUtils.TIMESTAMP_NTZ:
            return datetime.datetime
        elif type_info == VariantUtils.LONG_STR:
            return str
        raise PySparkValueError(errorClass="MALFORMED_VARIANT", messageParameters={})

    @classmethod
    def _to_json(cls, value: bytes, metadata: bytes, pos: int, zone_id: str) -> str:
        variant_type = cls._get_type(value, pos)
        if variant_type == dict:

            def handle_object(key_value_pos_list: List[Tuple[str, int]]) -> str:
                key_value_list = [
                    json.dumps(key) + ":" + cls._to_json(value, metadata, value_pos, zone_id)
                    for (key, value_pos) in key_value_pos_list
                ]
                return "{" + ",".join(key_value_list) + "}"

            return cls._handle_object(value, metadata, pos, handle_object)
        elif variant_type == array:

            def handle_array(value_pos_list: List[int]) -> str:
                value_list = [
                    cls._to_json(value, metadata, value_pos, zone_id)
                    for value_pos in value_pos_list
                ]
                return "[" + ",".join(value_list) + "]"

            return cls._handle_array(value, pos, handle_array)
        else:
            value = cls._get_scalar(variant_type, value, metadata, pos, zone_id)
            if value is None:
                return "null"
            if type(value) == bool:
                return "true" if value else "false"
            if type(value) == str:
                return json.dumps(value)
            if type(value) == bytes:
                # decoding simply converts byte array to string
                return '"' + base64.b64encode(value).decode("utf-8") + '"'
            if type(value) == datetime.date or type(value) == datetime.datetime:
                return '"' + str(value) + '"'
            return str(value)

    @classmethod
    def _to_python(cls, value: bytes, metadata: bytes, pos: int) -> Any:
        variant_type = cls._get_type(value, pos)
        if variant_type == dict:

            def handle_object(key_value_pos_list: List[Tuple[str, int]]) -> Dict[str, Any]:
                key_value_list = [
                    (key, cls._to_python(value, metadata, value_pos))
                    for (key, value_pos) in key_value_pos_list
                ]
                return dict(key_value_list)

            return cls._handle_object(value, metadata, pos, handle_object)
        elif variant_type == array:

            def handle_array(value_pos_list: List[int]) -> List[Any]:
                value_list = [
                    cls._to_python(value, metadata, value_pos) for value_pos in value_pos_list
                ]
                return value_list

            return cls._handle_array(value, pos, handle_array)
        else:
            return cls._get_scalar(variant_type, value, metadata, pos, zone_id="UTC")

    @classmethod
    def _get_scalar(
        cls, variant_type: Any, value: bytes, metadata: bytes, pos: int, zone_id: str
    ) -> Any:
        if isinstance(None, variant_type):
            return None
        elif variant_type == bool:
            return cls._get_boolean(value, pos)
        elif variant_type == int:
            return cls._get_long(value, pos)
        elif variant_type == str:
            return cls._get_string(value, pos)
        elif variant_type == float:
            return cls._get_double(value, pos)
        elif variant_type == decimal.Decimal:
            return cls._get_decimal(value, pos)
        elif variant_type == bytes:
            return cls._get_binary(value, pos)
        elif variant_type == datetime.date:
            return cls._get_date(value, pos)
        elif variant_type == datetime.datetime:
            return cls._get_timestamp(value, pos, zone_id)
        else:
            raise PySparkValueError(errorClass="MALFORMED_VARIANT", messageParameters={})

    @classmethod
    def _handle_object(
        cls, value: bytes, metadata: bytes, pos: int, func: Callable[[List[Tuple[str, int]]], Any]
    ) -> Any:
        """
        Parses the variant object at position `pos`.
        Calls `func` with a list of (key, value position) pairs of the object.
        """
        cls._check_index(pos, len(value))
        basic_type, type_info = cls._get_type_info(value, pos)
        if basic_type != VariantUtils.OBJECT:
            raise PySparkValueError(errorClass="MALFORMED_VARIANT", messageParameters={})
        large_size = ((type_info >> 4) & 0x1) != 0
        size_bytes = VariantUtils.U32_SIZE if large_size else 1
        num_fields = cls._read_long(value, pos + 1, size_bytes, signed=False)
        id_size = ((type_info >> 2) & 0x3) + 1
        offset_size = ((type_info) & 0x3) + 1
        id_start = pos + 1 + size_bytes
        offset_start = id_start + num_fields * id_size
        data_start = offset_start + (num_fields + 1) * offset_size

        key_value_pos_list = []
        for i in range(num_fields):
            id = cls._read_long(value, id_start + id_size * i, id_size, signed=False)
            offset = cls._read_long(
                value, offset_start + offset_size * i, offset_size, signed=False
            )
            value_pos = data_start + offset
            if metadata is not None:
                key_value_pos_list.append((cls._get_metadata_key(metadata, id), value_pos))
            else:
                key_value_pos_list.append(("", value_pos))
        return func(key_value_pos_list)

    @classmethod
    def _handle_array(cls, value: bytes, pos: int, func: Callable[[List[int]], Any]) -> Any:
        """
        Parses the variant array at position `pos`.
        Calls `func` with a list of element positions of the array.
        """
        cls._check_index(pos, len(value))
        basic_type, type_info = cls._get_type_info(value, pos)
        if basic_type != VariantUtils.ARRAY:
            raise PySparkValueError(errorClass="MALFORMED_VARIANT", messageParameters={})
        large_size = ((type_info >> 2) & 0x1) != 0
        size_bytes = VariantUtils.U32_SIZE if large_size else 1
        num_fields = cls._read_long(value, pos + 1, size_bytes, signed=False)
        offset_size = (type_info & 0x3) + 1
        offset_start = pos + 1 + size_bytes
        data_start = offset_start + (num_fields + 1) * offset_size

        value_pos_list = []
        for i in range(num_fields):
            offset = cls._read_long(
                value, offset_start + offset_size * i, offset_size, signed=False
            )
            element_pos = data_start + offset
            value_pos_list.append(element_pos)
        return func(value_pos_list)


class FieldEntry(NamedTuple):
    """
    Info about an object field
    """

    key: str
    id: int
    offset: int


class VariantBuilder:
    """
    A utility class for building VariantVal.
    """

    DEFAULT_SIZE_LIMIT = 16 * 1024 * 1024

    def __init__(self, size_limit: int = DEFAULT_SIZE_LIMIT):
        self.value = bytearray()
        self.dictionary = dict[str, int]()
        self.dictionary_keys = list[bytes]()
        self.size_limit = size_limit

    def build(self, json_str: str) -> Tuple[bytes, bytes]:
        parsed = json.loads(json_str, parse_float=self._handle_float)
        self._process_parsed_json(parsed)

        num_keys = len(self.dictionary_keys)
        dictionary_string_size = sum(len(key) for key in self.dictionary_keys)

        # Determine the number of bytes required per offset entry.
        # The largest offset is the one-past-the-end value, which is total string size. It's very
        # unlikely that the number of keys could be larger, but incorporate that into the
        # calculation in case of pathological data.
        max_size = max(dictionary_string_size, num_keys)
        if max_size > self.size_limit:
            raise PySparkValueError(errorClass="VARIANT_SIZE_LIMIT_EXCEEDED", messageParameters={})
        offset_size = self._get_integer_size(max_size)

        offset_start = 1 + offset_size
        string_start = offset_start + (num_keys + 1) * offset_size
        metadata_size = string_start + dictionary_string_size
        if metadata_size > self.size_limit:
            raise PySparkValueError(errorClass="VARIANT_SIZE_LIMIT_EXCEEDED", messageParameters={})

        metadata = bytearray()
        header_byte = VariantUtils.VERSION | ((offset_size - 1) << 6)
        metadata.extend(header_byte.to_bytes(1, byteorder="little"))
        metadata.extend(num_keys.to_bytes(offset_size, byteorder="little"))
        # write offsets
        current_offset = 0
        for key in self.dictionary_keys:
            metadata.extend(current_offset.to_bytes(offset_size, byteorder="little"))
            current_offset += len(key)
        metadata.extend(current_offset.to_bytes(offset_size, byteorder="little"))
        # write key data
        for key in self.dictionary_keys:
            metadata.extend(key)
        return (bytes(self.value), bytes(metadata))

    def _process_parsed_json(self, parsed: Any) -> None:
        if type(parsed) is dict:
            fields = list[FieldEntry]()
            start = len(self.value)
            for key, value in parsed.items():
                id = self._add_key(key)
                fields.append(FieldEntry(key, id, len(self.value) - start))
                self._process_parsed_json(value)
            self._finish_writing_object(start, fields)
        elif type(parsed) is list:
            offsets = []
            start = len(self.value)
            for elem in parsed:
                offsets.append(len(self.value) - start)
                self._process_parsed_json(elem)
            self._finish_writing_array(start, offsets)
        elif type(parsed) is str:
            self._append_string(parsed)
        elif type(parsed) is int:
            if not self._append_int(parsed):
                self._process_parsed_json(self._handle_float(str(parsed)))
        elif type(parsed) is float:
            self._append_float(parsed)
        elif type(parsed) is decimal.Decimal:
            self._append_decimal(parsed)
        elif type(parsed) is bool:
            self._append_boolean(parsed)
        elif parsed is None:
            self._append_null()
        else:
            raise PySparkValueError(errorClass="MALFORMED_VARIANT", messageParameters={})

    # Choose the smallest unsigned integer type that can store `value`. It must be within
    # [0, U24_MAX].
    def _get_integer_size(self, value: int) -> int:
        if value <= VariantUtils.U8_MAX:
            return 1
        if value <= VariantUtils.U16_MAX:
            return 2
        return VariantUtils.U24_SIZE

    def _check_capacity(self, additional: int) -> None:
        required = len(self.value) + additional
        if required > self.size_limit:
            raise PySparkValueError(errorClass="VARIANT_SIZE_LIMIT_EXCEEDED", messageParameters={})

    def _primitive_header(self, type: int) -> bytes:
        return bytes([(type << 2) | VariantUtils.PRIMITIVE])

    def _short_string_header(self, size: int) -> bytes:
        return bytes([size << 2 | VariantUtils.SHORT_STR])

    def _array_header(self, large_size: bool, offset_size: int) -> bytes:
        return bytes(
            [
                (
                    (large_size << (VariantUtils.BASIC_TYPE_BITS + 2))
                    | ((offset_size - 1) << VariantUtils.BASIC_TYPE_BITS)
                    | VariantUtils.ARRAY
                )
            ]
        )

    def _object_header(self, large_size: bool, id_size: int, offset_size: int) -> bytes:
        return bytes(
            [
                (
                    (large_size << (VariantUtils.BASIC_TYPE_BITS + 4))
                    | ((id_size - 1) << (VariantUtils.BASIC_TYPE_BITS + 2))
                    | ((offset_size - 1) << VariantUtils.BASIC_TYPE_BITS)
                    | VariantUtils.OBJECT
                )
            ]
        )

    # Add a key to the variant dictionary. If the key already exists, the dictionary is
    # not modified. In either case, return the id of the key.
    def _add_key(self, key: str) -> int:
        if key in self.dictionary:
            return self.dictionary[key]
        id = len(self.dictionary_keys)
        self.dictionary[key] = id
        self.dictionary_keys.append(key.encode("utf-8"))
        return id

    def _handle_float(self, num_str: str) -> Any:
        # a float can be a decimal if it only contains digits, '-', or '-'.
        if all([ch.isdecimal() or ch == "-" or ch == "." for ch in num_str]):
            dec = decimal.Decimal(num_str)
            precision = len(dec.as_tuple().digits)
            scale = -int(dec.as_tuple().exponent)

            if (
                scale <= VariantUtils.MAX_DECIMAL16_PRECISION
                and precision <= VariantUtils.MAX_DECIMAL16_PRECISION
            ):
                return dec
        return float(num_str)

    def _append_boolean(self, b: bool) -> None:
        self._check_capacity(1)
        self.value.extend(self._primitive_header(VariantUtils.TRUE if b else VariantUtils.FALSE))

    def _append_null(self) -> None:
        self._check_capacity(1)
        self.value.extend(self._primitive_header(VariantUtils.NULL))

    def _append_string(self, s: str) -> None:
        text = s.encode("utf-8")
        long_str = len(text) > VariantUtils.MAX_SHORT_STR_SIZE
        additional = (1 + VariantUtils.U32_SIZE) if long_str else 1
        self._check_capacity(additional + len(text))
        if long_str:
            self.value.extend(self._primitive_header(VariantUtils.LONG_STR))
            self.value.extend(len(text).to_bytes(VariantUtils.U32_SIZE, byteorder="little"))
        else:
            self.value.extend(self._short_string_header(len(text)))
        self.value.extend(text)

    def _append_int(self, i: int) -> bool:
        self._check_capacity(1 + 8)
        if i >= VariantUtils.I8_MIN and i <= VariantUtils.I8_MAX:
            self.value.extend(self._primitive_header(VariantUtils.INT1))
            self.value.extend(i.to_bytes(1, byteorder="little", signed=True))
        elif i >= VariantUtils.I16_MIN and i <= VariantUtils.I16_MAX:
            self.value.extend(self._primitive_header(VariantUtils.INT2))
            self.value.extend(i.to_bytes(2, byteorder="little", signed=True))
        elif i >= VariantUtils.I32_MIN and i <= VariantUtils.I32_MAX:
            self.value.extend(self._primitive_header(VariantUtils.INT4))
            self.value.extend(i.to_bytes(4, byteorder="little", signed=True))
        elif i >= VariantUtils.I64_MIN and i <= VariantUtils.I64_MAX:
            self.value.extend(self._primitive_header(VariantUtils.INT8))
            self.value.extend(i.to_bytes(8, byteorder="little", signed=True))
        else:
            return False
        return True

    # Append a decimal value to the variant builder. The caller should guarantee that its precision
    # and scale fit into `MAX_DECIMAL16_PRECISION`.
    def _append_decimal(self, d: decimal.Decimal) -> None:
        self._check_capacity(2 + 16)
        precision = len(d.as_tuple().digits)
        scale = -int(d.as_tuple().exponent)
        unscaled = int("".join(map(str, d.as_tuple().digits)))
        unscaled = -unscaled if d < 0 else unscaled
        if (
            scale <= VariantUtils.MAX_DECIMAL4_PRECISION
            and precision <= VariantUtils.MAX_DECIMAL4_PRECISION
        ):
            self.value.extend(self._primitive_header(VariantUtils.DECIMAL4))
            self.value.extend(scale.to_bytes(1, byteorder="little"))
            self.value.extend(unscaled.to_bytes(4, byteorder="little", signed=True))
        elif (
            scale <= VariantUtils.MAX_DECIMAL8_PRECISION
            and precision <= VariantUtils.MAX_DECIMAL8_PRECISION
        ):
            self.value.extend(self._primitive_header(VariantUtils.DECIMAL8))
            self.value.extend(scale.to_bytes(1, byteorder="little"))
            self.value.extend(unscaled.to_bytes(8, byteorder="little", signed=True))
        else:
            assert (
                scale <= VariantUtils.MAX_DECIMAL16_PRECISION
                and precision <= VariantUtils.MAX_DECIMAL16_PRECISION
            )
            self.value.extend(self._primitive_header(VariantUtils.DECIMAL16))
            self.value.extend(scale.to_bytes(1, byteorder="little"))
            self.value.extend(unscaled.to_bytes(16, byteorder="little", signed=True))

    def _append_float(self, f: float) -> None:
        self._check_capacity(1 + 8)
        self.value.extend(self._primitive_header(VariantUtils.DOUBLE))
        self.value.extend(struct.pack("<d", f))

    # Finish writing a variant array after all of its elements have already been written.
    def _finish_writing_array(self, start: int, offsets: List[int]) -> None:
        data_size = len(self.value) - start
        num_offsets = len(offsets)
        large_size = num_offsets > VariantUtils.U8_MAX
        size_bytes = VariantUtils.U32_SIZE if large_size else 1
        offset_size = self._get_integer_size(data_size)
        # The space for header byte, object size, and offset list.
        header_size = 1 + size_bytes + (num_offsets + 1) * offset_size
        self._check_capacity(header_size)
        self.value.extend(bytearray(header_size))
        # Shift the just-written element data to make room for the header section.
        self.value[start + header_size :] = bytes(self.value[start : start + data_size])
        # Write the header byte, num offsets
        offset_start = start + 1 + size_bytes
        self.value[start : start + 1] = self._array_header(large_size, offset_size)
        self.value[start + 1 : offset_start] = num_offsets.to_bytes(size_bytes, byteorder="little")
        # write offset list
        offset_list = bytearray()
        for offset in offsets:
            offset_list.extend(offset.to_bytes(offset_size, byteorder="little"))
        offset_list.extend(data_size.to_bytes(offset_size, byteorder="little"))
        self.value[offset_start : offset_start + len(offset_list)] = offset_list

    # Finish writing a variant object after all of its fields have already been written.
    def _finish_writing_object(self, start: int, fields: List[FieldEntry]) -> None:
        num_fields = len(fields)
        # object fields are from a python dictionary, so keys are already distinct
        fields.sort(key=lambda f: f.key)
        max_id = 0
        for field in fields:
            max_id = max(max_id, field.id)

        data_size = len(self.value) - start
        large_size = num_fields > VariantUtils.U8_MAX
        size_bytes = VariantUtils.U32_SIZE if large_size else 1
        id_size = self._get_integer_size(max_id)
        offset_size = self._get_integer_size(data_size)
        # The space for header byte, object size, id list, and offset list.
        header_size = 1 + size_bytes + num_fields * id_size + (num_fields + 1) * offset_size
        self._check_capacity(header_size)
        self.value.extend(bytearray(header_size))
        # Shift the just-written field data to make room for the object header section.
        self.value[start + header_size :] = self.value[start : start + data_size]
        # Write the header byte, num fields, id list, offset list
        self.value[start : start + 1] = self._object_header(large_size, id_size, offset_size)
        self.value[start + 1 : start + 1 + size_bytes] = num_fields.to_bytes(
            size_bytes, byteorder="little"
        )
        id_start = start + 1 + size_bytes
        offset_start = id_start + num_fields * id_size
        id_list = bytearray()
        offset_list = bytearray()
        for field in fields:
            id_list.extend(field.id.to_bytes(id_size, byteorder="little"))
            offset_list.extend(field.offset.to_bytes(offset_size, byteorder="little"))
        offset_list.extend(data_size.to_bytes(offset_size, byteorder="little"))
        self.value[id_start : id_start + len(id_list)] = id_list
        self.value[offset_start : offset_start + len(offset_list)] = offset_list