# Lint as: python3
# Copyright 2020 Google Inc. All rights reserved.
#
# Licensed 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.
"""Implementation of FlexBuffers binary format.

For more info check https://google.github.io/flatbuffers/flexbuffers.html and
corresponding C++ implementation at
https://github.com/google/flatbuffers/blob/master/include/flatbuffers/flexbuffers.h
"""

# pylint: disable=invalid-name
# TODO(dkovalev): Add type hints everywhere, so tools like pytypes could work.

import array
import contextlib
import enum
import struct

__all__ = ("Type", "Builder", "GetRoot", "Dumps", "Loads")


class BitWidth(enum.IntEnum):
    """Supported bit widths of value types.

    These are used in the lower 2 bits of a type field to determine the size of
    the elements (and or size field) of the item pointed to (e.g. vector).
    """

    W8 = 0  # 2^0 = 1 byte
    W16 = 1  # 2^1 = 2 bytes
    W32 = 2  # 2^2 = 4 bytes
    W64 = 3  # 2^3 = 8 bytes

    @staticmethod
    def U(value):
        """Returns the minimum `BitWidth` to encode unsigned integer value."""
        assert value >= 0

        if value < (1 << 8):
            return BitWidth.W8
        elif value < (1 << 16):
            return BitWidth.W16
        elif value < (1 << 32):
            return BitWidth.W32
        elif value < (1 << 64):
            return BitWidth.W64
        else:
            raise ValueError("value is too big to encode: %s" % value)

    @staticmethod
    def I(value):
        """Returns the minimum `BitWidth` to encode signed integer value."""
        # -2^(n-1) <=     value < 2^(n-1)
        # -2^n     <= 2 * value < 2^n
        # 2 * value < 2^n, when value >= 0 or 2 * (-value) <= 2^n, when value < 0
        # 2 * value < 2^n, when value >= 0 or 2 * (-value) - 1 < 2^n, when value < 0
        #
        # if value >= 0:
        #   return BitWidth.U(2 * value)
        # else:
        #   return BitWidth.U(2 * (-value) - 1)  # ~x = -x - 1
        value *= 2
        return BitWidth.U(value if value >= 0 else ~value)

    @staticmethod
    def F(value):
        """Returns the `BitWidth` to encode floating point value."""
        if struct.unpack("f", struct.pack("f", value))[0] == value:
            return BitWidth.W32
        return BitWidth.W64

    @staticmethod
    def B(byte_width):
        return {1: BitWidth.W8, 2: BitWidth.W16, 4: BitWidth.W32, 8: BitWidth.W64}[
            byte_width
        ]


I = {1: "b", 2: "h", 4: "i", 8: "q"}  # Integer formats
U = {1: "B", 2: "H", 4: "I", 8: "Q"}  # Unsigned integer formats
F = {4: "f", 8: "d"}  # Floating point formats


def _Unpack(fmt, buf):
    return struct.unpack(fmt[len(buf)], buf)[0]


def _UnpackVector(fmt, buf, length):
    byte_width = len(buf) // length
    return struct.unpack("%d%s" % (length, fmt[byte_width]), buf)


def _Pack(fmt, value, byte_width):
    return struct.pack(fmt[byte_width], value)


def _PackVector(fmt, values, byte_width):
    return struct.pack("%d%s" % (len(values), fmt[byte_width]), *values)


def _Mutate(fmt, buf, value, byte_width, value_bit_width):
    if (1 << value_bit_width) <= byte_width:
        buf[:byte_width] = _Pack(fmt, value, byte_width)
        return True
    return False


# Computes how many bytes you'd have to pad to be able to write an
# "scalar_size" scalar if the buffer had grown to "buf_size",
# "scalar_size" is a power of two.
def _PaddingBytes(buf_size, scalar_size):
    # ((buf_size + (scalar_size - 1)) // scalar_size) * scalar_size - buf_size
    return -buf_size & (scalar_size - 1)


def _ShiftSlice(s, offset, length):
    start = offset + (0 if s.start is None else s.start)
    stop = offset + (length if s.stop is None else s.stop)
    return slice(start, stop, s.step)


# https://en.cppreference.com/w/cpp/algorithm/lower_bound
def _LowerBound(values, value, pred):
    """Implementation of C++ std::lower_bound() algorithm."""
    first, last = 0, len(values)
    count = last - first
    while count > 0:
        i = first
        step = count // 2
        i += step
        if pred(values[i], value):
            i += 1
            first = i
            count -= step + 1
        else:
            count = step
    return first


# https://en.cppreference.com/w/cpp/algorithm/binary_search
def _BinarySearch(values, value, pred=lambda x, y: x < y):
    """Implementation of C++ std::binary_search() algorithm."""
    index = _LowerBound(values, value, pred)
    if index != len(values) and not pred(value, values[index]):
        return index
    return -1


class Type(enum.IntEnum):
    """Supported types of encoded data.

    These are used as the upper 6 bits of a type field to indicate the actual
    type.
    """

    NULL = 0
    INT = 1
    UINT = 2
    FLOAT = 3
    # Types above stored inline, types below store an offset.
    KEY = 4
    STRING = 5
    INDIRECT_INT = 6
    INDIRECT_UINT = 7
    INDIRECT_FLOAT = 8
    MAP = 9
    VECTOR = 10  # Untyped.

    VECTOR_INT = 11  # Typed any size (stores no type table).
    VECTOR_UINT = 12
    VECTOR_FLOAT = 13
    VECTOR_KEY = 14
    # DEPRECATED, use VECTOR or VECTOR_KEY instead.
    # Read test.cpp/FlexBuffersDeprecatedTest() for details on why.
    VECTOR_STRING_DEPRECATED = 15

    VECTOR_INT2 = 16  # Typed tuple (no type table, no size field).
    VECTOR_UINT2 = 17
    VECTOR_FLOAT2 = 18
    VECTOR_INT3 = 19  # Typed triple (no type table, no size field).
    VECTOR_UINT3 = 20
    VECTOR_FLOAT3 = 21
    VECTOR_INT4 = 22  # Typed quad (no type table, no size field).
    VECTOR_UINT4 = 23
    VECTOR_FLOAT4 = 24

    BLOB = 25
    BOOL = 26
    VECTOR_BOOL = 36  # To do the same type of conversion of type to vector type

    @staticmethod
    def Pack(type_, bit_width):
        return (int(type_) << 2) | bit_width

    @staticmethod
    def Unpack(packed_type):
        return 1 << (packed_type & 0b11), Type(packed_type >> 2)

    @staticmethod
    def IsInline(type_):
        return type_ <= Type.FLOAT or type_ == Type.BOOL

    @staticmethod
    def IsTypedVector(type_):
        return (
            Type.VECTOR_INT <= type_ <= Type.VECTOR_STRING_DEPRECATED
            or type_ == Type.VECTOR_BOOL
        )

    @staticmethod
    def IsTypedVectorElementType(type_):
        return Type.INT <= type_ <= Type.STRING or type_ == Type.BOOL

    @staticmethod
    def ToTypedVectorElementType(type_):
        if not Type.IsTypedVector(type_):
            raise ValueError("must be typed vector type")

        return Type(type_ - Type.VECTOR_INT + Type.INT)

    @staticmethod
    def IsFixedTypedVector(type_):
        return Type.VECTOR_INT2 <= type_ <= Type.VECTOR_FLOAT4

    @staticmethod
    def IsFixedTypedVectorElementType(type_):
        return Type.INT <= type_ <= Type.FLOAT

    @staticmethod
    def ToFixedTypedVectorElementType(type_):
        if not Type.IsFixedTypedVector(type_):
            raise ValueError("must be fixed typed vector type")

        # 3 types each, starting from length 2.
        fixed_type = type_ - Type.VECTOR_INT2
        return Type(fixed_type % 3 + Type.INT), fixed_type // 3 + 2

    @staticmethod
    def ToTypedVector(element_type, fixed_len=0):
        """Converts element type to corresponding vector type.

        Args:
          element_type: vector element type
          fixed_len: number of elements: 0 for typed vector; 2, 3, or 4 for fixed
            typed vector.

        Returns:
          Typed vector type or fixed typed vector type.
        """
        if fixed_len == 0:
            if not Type.IsTypedVectorElementType(element_type):
                raise ValueError("must be typed vector element type")
        else:
            if not Type.IsFixedTypedVectorElementType(element_type):
                raise ValueError("must be fixed typed vector element type")

        offset = element_type - Type.INT
        if fixed_len == 0:
            return Type(offset + Type.VECTOR_INT)  # TypedVector
        elif fixed_len == 2:
            return Type(offset + Type.VECTOR_INT2)  # FixedTypedVector
        elif fixed_len == 3:
            return Type(offset + Type.VECTOR_INT3)  # FixedTypedVector
        elif fixed_len == 4:
            return Type(offset + Type.VECTOR_INT4)  # FixedTypedVector
        else:
            raise ValueError("unsupported fixed_len: %s" % fixed_len)


class Buf:
    """Class to access underlying buffer object starting from the given offset."""

    def __init__(self, buf, offset):
        self._buf = buf
        self._offset = offset if offset >= 0 else len(buf) + offset
        self._length = len(buf) - self._offset

    def __getitem__(self, key):
        if isinstance(key, slice):
            return self._buf[_ShiftSlice(key, self._offset, self._length)]
        elif isinstance(key, int):
            return self._buf[self._offset + key]
        else:
            raise TypeError("invalid key type")

    def __setitem__(self, key, value):
        if isinstance(key, slice):
            self._buf[_ShiftSlice(key, self._offset, self._length)] = value
        elif isinstance(key, int):
            self._buf[self._offset + key] = key
        else:
            raise TypeError("invalid key type")

    def __repr__(self):
        return "buf[%d:]" % self._offset

    def Find(self, sub):
        """Returns the lowest index where the sub subsequence is found."""
        return self._buf[self._offset :].find(sub)

    def Slice(self, offset):
        """Returns new `Buf` which starts from the given offset."""
        return Buf(self._buf, self._offset + offset)

    def Indirect(self, offset, byte_width):
        """Return new `Buf` based on the encoded offset (indirect encoding)."""
        return self.Slice(offset - _Unpack(U, self[offset : offset + byte_width]))


class Object:
    """Base class for all non-trivial data accessors."""

    __slots__ = "_buf", "_byte_width"

    def __init__(self, buf, byte_width):
        self._buf = buf
        self._byte_width = byte_width

    @property
    def ByteWidth(self):
        return self._byte_width


class Sized(Object):
    """Base class for all data accessors which need to read encoded size."""

    __slots__ = ("_size",)

    def __init__(self, buf, byte_width, size=0):
        super().__init__(buf, byte_width)
        if size == 0:
            self._size = _Unpack(U, self.SizeBytes)
        else:
            self._size = size

    @property
    def SizeBytes(self):
        return self._buf[-self._byte_width : 0]

    def __len__(self):
        return self._size


class Blob(Sized):
    """Data accessor for the encoded blob bytes."""

    __slots__ = ()

    @property
    def Bytes(self):
        return self._buf[0 : len(self)]

    def __repr__(self):
        return "Blob(%s, size=%d)" % (self._buf, len(self))


class String(Sized):
    """Data accessor for the encoded string bytes."""

    __slots__ = ()

    @property
    def Bytes(self):
        return self._buf[0 : len(self)]

    def Mutate(self, value):
        """Mutates underlying string bytes in place.

        Args:
          value: New string to replace the existing one. New string must have less
            or equal UTF-8-encoded bytes than the existing one to successfully
            mutate underlying byte buffer.

        Returns:
          Whether the value was mutated or not.
        """
        encoded = value.encode("utf-8")
        n = len(encoded)
        if n <= len(self):
            self._buf[-self._byte_width : 0] = _Pack(U, n, self._byte_width)
            self._buf[0:n] = encoded
            self._buf[n : len(self)] = bytearray(len(self) - n)
            return True
        return False

    def __str__(self):
        return self.Bytes.decode("utf-8")

    def __repr__(self):
        return "String(%s, size=%d)" % (self._buf, len(self))


class Key(Object):
    """Data accessor for the encoded key bytes."""

    __slots__ = ()

    def __init__(self, buf, byte_width):
        assert byte_width == 1
        super().__init__(buf, byte_width)

    @property
    def Bytes(self):
        return self._buf[0 : len(self)]

    def __len__(self):
        return self._buf.Find(0)

    def __str__(self):
        return self.Bytes.decode("ascii")

    def __repr__(self):
        return "Key(%s, size=%d)" % (self._buf, len(self))


class Vector(Sized):
    """Data accessor for the encoded vector bytes."""

    __slots__ = ()

    def __getitem__(self, index):
        if index < 0 or index >= len(self):
            raise IndexError(
                "vector index %s is out of [0, %d) range" % (index, len(self))
            )

        packed_type = self._buf[len(self) * self._byte_width + index]
        buf = self._buf.Slice(index * self._byte_width)
        return Ref.PackedType(buf, self._byte_width, packed_type)

    @property
    def Value(self):
        """Returns the underlying encoded data as a list object."""
        return [e.Value for e in self]

    def __repr__(self):
        return "Vector(%s, byte_width=%d, size=%d)" % (
            self._buf,
            self._byte_width,
            self._size,
        )


class TypedVector(Sized):
    """Data accessor for the encoded typed vector or fixed typed vector bytes."""

    __slots__ = "_element_type", "_size"

    def __init__(self, buf, byte_width, element_type, size=0):
        super().__init__(buf, byte_width, size)

        if element_type == Type.STRING:
            # These can't be accessed as strings, since we don't know the bit-width
            # of the size field, see the declaration of
            # FBT_VECTOR_STRING_DEPRECATED above for details.
            # We change the type here to be keys, which are a subtype of strings,
            # and will ignore the size field. This will truncate strings with
            # embedded nulls.
            element_type = Type.KEY

        self._element_type = element_type

    @property
    def Bytes(self):
        return self._buf[: self._byte_width * len(self)]

    @property
    def ElementType(self):
        return self._element_type

    def __getitem__(self, index):
        if index < 0 or index >= len(self):
            raise IndexError(
                "vector index %s is out of [0, %d) range" % (index, len(self))
            )

        buf = self._buf.Slice(index * self._byte_width)
        return Ref(buf, self._byte_width, 1, self._element_type)

    @property
    def Value(self):
        """Returns underlying data as list object."""
        if not self:
            return []

        if self._element_type is Type.BOOL:
            return [bool(e) for e in _UnpackVector(U, self.Bytes, len(self))]
        elif self._element_type is Type.INT:
            return list(_UnpackVector(I, self.Bytes, len(self)))
        elif self._element_type is Type.UINT:
            return list(_UnpackVector(U, self.Bytes, len(self)))
        elif self._element_type is Type.FLOAT:
            return list(_UnpackVector(F, self.Bytes, len(self)))
        elif self._element_type is Type.KEY:
            return [e.AsKey for e in self]
        elif self._element_type is Type.STRING:
            return [e.AsString for e in self]
        else:
            raise TypeError("unsupported element_type: %s" % self._element_type)

    def __repr__(self):
        return "TypedVector(%s, byte_width=%d, element_type=%s, size=%d)" % (
            self._buf,
            self._byte_width,
            self._element_type,
            self._size,
        )


class Map(Vector):
    """Data accessor for the encoded map bytes."""

    @staticmethod
    def CompareKeys(a, b):
        if isinstance(a, Ref):
            a = a.AsKeyBytes
        if isinstance(b, Ref):
            b = b.AsKeyBytes
        return a < b

    def __getitem__(self, key):
        if isinstance(key, int):
            return super().__getitem__(key)

        index = _BinarySearch(self.Keys, key.encode("ascii"), self.CompareKeys)
        if index != -1:
            return super().__getitem__(index)

        raise KeyError(key)

    @property
    def Keys(self):
        byte_width = _Unpack(U, self._buf[-2 * self._byte_width : -self._byte_width])
        buf = self._buf.Indirect(-3 * self._byte_width, self._byte_width)
        return TypedVector(buf, byte_width, Type.KEY)

    @property
    def Values(self):
        return Vector(self._buf, self._byte_width)

    @property
    def Value(self):
        return {k.Value: v.Value for k, v in zip(self.Keys, self.Values)}

    def __repr__(self):
        return "Map(%s, size=%d)" % (self._buf, len(self))


class Ref:
    """Data accessor for the encoded data bytes."""

    __slots__ = "_buf", "_parent_width", "_byte_width", "_type"

    @staticmethod
    def PackedType(buf, parent_width, packed_type):
        byte_width, type_ = Type.Unpack(packed_type)
        return Ref(buf, parent_width, byte_width, type_)

    def __init__(self, buf, parent_width, byte_width, type_):
        self._buf = buf
        self._parent_width = parent_width
        self._byte_width = byte_width
        self._type = type_

    def __repr__(self):
        return "Ref(%s, parent_width=%d, byte_width=%d, type_=%s)" % (
            self._buf,
            self._parent_width,
            self._byte_width,
            self._type,
        )

    @property
    def _Bytes(self):
        return self._buf[: self._parent_width]

    def _ConvertError(self, target_type):
        raise TypeError("cannot convert %s to %s" % (self._type, target_type))

    def _Indirect(self):
        return self._buf.Indirect(0, self._parent_width)

    @property
    def IsNull(self):
        return self._type is Type.NULL

    @property
    def IsBool(self):
        return self._type is Type.BOOL

    @property
    def AsBool(self):
        if self._type is Type.BOOL:
            return bool(_Unpack(U, self._Bytes))
        else:
            return self.AsInt != 0

    def MutateBool(self, value):
        """Mutates underlying boolean value bytes in place.

        Args:
          value: New boolean value.

        Returns:
          Whether the value was mutated or not.
        """
        return self.IsBool and _Mutate(
            U, self._buf, value, self._parent_width, BitWidth.W8
        )

    @property
    def IsNumeric(self):
        return self.IsInt or self.IsFloat

    @property
    def IsInt(self):
        return self._type in (
            Type.INT,
            Type.INDIRECT_INT,
            Type.UINT,
            Type.INDIRECT_UINT,
        )

    @property
    def AsInt(self):
        """Returns current reference as integer value."""
        if self.IsNull:
            return 0
        elif self.IsBool:
            return int(self.AsBool)
        elif self._type is Type.INT:
            return _Unpack(I, self._Bytes)
        elif self._type is Type.INDIRECT_INT:
            return _Unpack(I, self._Indirect()[: self._byte_width])
        if self._type is Type.UINT:
            return _Unpack(U, self._Bytes)
        elif self._type is Type.INDIRECT_UINT:
            return _Unpack(U, self._Indirect()[: self._byte_width])
        elif self.IsString:
            return len(self.AsString)
        elif self.IsKey:
            return len(self.AsKey)
        elif self.IsBlob:
            return len(self.AsBlob)
        elif self.IsVector:
            return len(self.AsVector)
        elif self.IsTypedVector:
            return len(self.AsTypedVector)
        elif self.IsFixedTypedVector:
            return len(self.AsFixedTypedVector)
        else:
            raise self._ConvertError(Type.INT)

    def MutateInt(self, value):
        """Mutates underlying integer value bytes in place.

        Args:
          value: New integer value. It must fit to the byte size of the existing
            encoded value.

        Returns:
          Whether the value was mutated or not.
        """
        if self._type is Type.INT:
            return _Mutate(I, self._buf, value, self._parent_width, BitWidth.I(value))
        elif self._type is Type.INDIRECT_INT:
            return _Mutate(
                I, self._Indirect(), value, self._byte_width, BitWidth.I(value)
            )
        elif self._type is Type.UINT:
            return _Mutate(U, self._buf, value, self._parent_width, BitWidth.U(value))
        elif self._type is Type.INDIRECT_UINT:
            return _Mutate(
                U, self._Indirect(), value, self._byte_width, BitWidth.U(value)
            )
        else:
            return False

    @property
    def IsFloat(self):
        return self._type in (Type.FLOAT, Type.INDIRECT_FLOAT)

    @property
    def AsFloat(self):
        """Returns current reference as floating point value."""
        if self.IsNull:
            return 0.0
        elif self.IsBool:
            return float(self.AsBool)
        elif self.IsInt:
            return float(self.AsInt)
        elif self._type is Type.FLOAT:
            return _Unpack(F, self._Bytes)
        elif self._type is Type.INDIRECT_FLOAT:
            return _Unpack(F, self._Indirect()[: self._byte_width])
        elif self.IsString:
            return float(self.AsString)
        elif self.IsVector:
            return float(len(self.AsVector))
        elif self.IsTypedVector():
            return float(len(self.AsTypedVector))
        elif self.IsFixedTypedVector():
            return float(len(self.FixedTypedVector))
        else:
            raise self._ConvertError(Type.FLOAT)

    def MutateFloat(self, value):
        """Mutates underlying floating point value bytes in place.

        Args:
          value: New float value. It must fit to the byte size of the existing
            encoded value.

        Returns:
          Whether the value was mutated or not.
        """
        if self._type is Type.FLOAT:
            return _Mutate(
                F, self._buf, value, self._parent_width, BitWidth.B(self._parent_width)
            )
        elif self._type is Type.INDIRECT_FLOAT:
            return _Mutate(
                F,
                self._Indirect(),
                value,
                self._byte_width,
                BitWidth.B(self._byte_width),
            )
        else:
            return False

    @property
    def IsKey(self):
        return self._type is Type.KEY

    @property
    def AsKeyBytes(self):
        if self.IsKey:
            return Key(self._Indirect(), self._byte_width).Bytes
        else:
            raise self._ConvertError(Type.KEY)

    @property
    def AsKey(self):
        if self.IsKey:
            return str(Key(self._Indirect(), self._byte_width))
        else:
            raise self._ConvertError(Type.KEY)

    @property
    def IsString(self):
        return self._type is Type.STRING

    @property
    def AsString(self):
        if self.IsString:
            return str(String(self._Indirect(), self._byte_width))
        elif self.IsKey:
            return self.AsKey
        else:
            raise self._ConvertError(Type.STRING)

    def MutateString(self, value):
        return String(self._Indirect(), self._byte_width).Mutate(value)

    @property
    def IsBlob(self):
        return self._type is Type.BLOB

    @property
    def AsBlob(self):
        if self.IsBlob:
            return Blob(self._Indirect(), self._byte_width).Bytes
        else:
            raise self._ConvertError(Type.BLOB)

    @property
    def IsAnyVector(self):
        return self.IsVector or self.IsTypedVector or self.IsFixedTypedVector()

    @property
    def IsVector(self):
        return self._type in (Type.VECTOR, Type.MAP)

    @property
    def AsVector(self):
        if self.IsVector:
            return Vector(self._Indirect(), self._byte_width)
        else:
            raise self._ConvertError(Type.VECTOR)

    @property
    def IsTypedVector(self):
        return Type.IsTypedVector(self._type)

    @property
    def AsTypedVector(self):
        if self.IsTypedVector:
            return TypedVector(
                self._Indirect(),
                self._byte_width,
                Type.ToTypedVectorElementType(self._type),
            )
        else:
            raise self._ConvertError("TYPED_VECTOR")

    @property
    def IsFixedTypedVector(self):
        return Type.IsFixedTypedVector(self._type)

    @property
    def AsFixedTypedVector(self):
        if self.IsFixedTypedVector:
            element_type, size = Type.ToFixedTypedVectorElementType(self._type)
            return TypedVector(self._Indirect(), self._byte_width, element_type, size)
        else:
            raise self._ConvertError("FIXED_TYPED_VECTOR")

    @property
    def IsMap(self):
        return self._type is Type.MAP

    @property
    def AsMap(self):
        if self.IsMap:
            return Map(self._Indirect(), self._byte_width)
        else:
            raise self._ConvertError(Type.MAP)

    @property
    def Value(self):
        """Converts current reference to value of corresponding type.

        This is equivalent to calling `AsInt` for integer values, `AsFloat` for
        floating point values, etc.

        Returns:
          Value of corresponding type.
        """
        if self.IsNull:
            return None
        elif self.IsBool:
            return self.AsBool
        elif self.IsInt:
            return self.AsInt
        elif self.IsFloat:
            return self.AsFloat
        elif self.IsString:
            return self.AsString
        elif self.IsKey:
            return self.AsKey
        elif self.IsBlob:
            return self.AsBlob
        elif self.IsMap:
            return self.AsMap.Value
        elif self.IsVector:
            return self.AsVector.Value
        elif self.IsTypedVector:
            return self.AsTypedVector.Value
        elif self.IsFixedTypedVector:
            return self.AsFixedTypedVector.Value
        else:
            raise TypeError("cannot convert %r to value" % self)


def _IsIterable(obj):
    try:
        iter(obj)
        return True
    except TypeError:
        return False


class Value:
    """Class to represent given value during the encoding process."""

    @staticmethod
    def Null():
        return Value(0, Type.NULL, BitWidth.W8)

    @staticmethod
    def Bool(value):
        return Value(value, Type.BOOL, BitWidth.W8)

    @staticmethod
    def Int(value, bit_width):
        return Value(value, Type.INT, bit_width)

    @staticmethod
    def UInt(value, bit_width):
        return Value(value, Type.UINT, bit_width)

    @staticmethod
    def Float(value, bit_width):
        return Value(value, Type.FLOAT, bit_width)

    @staticmethod
    def Key(offset):
        return Value(offset, Type.KEY, BitWidth.W8)

    def __init__(self, value, type_, min_bit_width):
        self._value = value
        self._type = type_

        # For scalars: of itself, for vector: of its elements, for string: length.
        self._min_bit_width = min_bit_width

    @property
    def Value(self):
        return self._value

    @property
    def Type(self):
        return self._type

    @property
    def MinBitWidth(self):
        return self._min_bit_width

    def StoredPackedType(self, parent_bit_width=BitWidth.W8):
        return Type.Pack(self._type, self.StoredWidth(parent_bit_width))

    # We have an absolute offset, but want to store a relative offset
    # elem_index elements beyond the current buffer end. Since whether
    # the relative offset fits in a certain byte_width depends on
    # the size of the elements before it (and their alignment), we have
    # to test for each size in turn.
    def ElemWidth(self, buf_size, elem_index=0):
        if Type.IsInline(self._type):
            return self._min_bit_width
        for byte_width in 1, 2, 4, 8:
            offset_loc = (
                buf_size + _PaddingBytes(buf_size, byte_width) + elem_index * byte_width
            )
            bit_width = BitWidth.U(offset_loc - self._value)
            if byte_width == (1 << bit_width):
                return bit_width
        raise ValueError("relative offset is too big")

    def StoredWidth(self, parent_bit_width=BitWidth.W8):
        if Type.IsInline(self._type):
            return max(self._min_bit_width, parent_bit_width)
        return self._min_bit_width

    def __repr__(self):
        return "Value(%s, %s, %s)" % (self._value, self._type, self._min_bit_width)

    def __str__(self):
        return str(self._value)


def InMap(func):
    def wrapper(self, *args, **kwargs):
        if isinstance(args[0], str):
            self.Key(args[0])
            func(self, *args[1:], **kwargs)
        else:
            func(self, *args, **kwargs)

    return wrapper


def InMapForString(func):
    def wrapper(self, *args):
        if len(args) == 1:
            func(self, args[0])
        elif len(args) == 2:
            self.Key(args[0])
            func(self, args[1])
        else:
            raise ValueError("invalid number of arguments")

    return wrapper


class Pool:
    """Collection of (data, offset) pairs sorted by data for quick access."""

    def __init__(self):
        self._pool = []  # sorted list of (data, offset) tuples

    def FindOrInsert(self, data, offset):
        do = data, offset
        index = _BinarySearch(self._pool, do, lambda a, b: a[0] < b[0])
        if index != -1:
            _, offset = self._pool[index]
            return offset
        self._pool.insert(index, do)
        return None

    def Clear(self):
        self._pool = []

    @property
    def Elements(self):
        return [data for data, _ in self._pool]


class Builder:
    """Helper class to encode structural data into flexbuffers format."""

    def __init__(
        self, share_strings=False, share_keys=True, force_min_bit_width=BitWidth.W8
    ):
        self._share_strings = share_strings
        self._share_keys = share_keys
        self._force_min_bit_width = force_min_bit_width

        self._string_pool = Pool()
        self._key_pool = Pool()

        self._finished = False
        self._buf = bytearray()
        self._stack = []

    def __len__(self):
        return len(self._buf)

    @property
    def StringPool(self):
        return self._string_pool

    @property
    def KeyPool(self):
        return self._key_pool

    def Clear(self):
        self._string_pool.Clear()
        self._key_pool.Clear()
        self._finished = False
        self._buf = bytearray()
        self._stack = []

    def Finish(self):
        """Finishes encoding process and returns underlying buffer."""
        if self._finished:
            raise RuntimeError("builder has been already finished")

        # If you hit this exception, you likely have objects that were never
        # included in a parent. You need to have exactly one root to finish a
        # buffer. Check your Start/End calls are matched, and all objects are inside
        # some other object.
        if len(self._stack) != 1:
            raise RuntimeError("internal stack size must be one")

        value = self._stack[0]
        byte_width = self._Align(value.ElemWidth(len(self._buf)))
        self._WriteAny(value, byte_width=byte_width)  # Root value
        self._Write(U, value.StoredPackedType(), byte_width=1)  # Root type
        self._Write(U, byte_width, byte_width=1)  # Root size

        self.finished = True
        return self._buf

    def _ReadKey(self, offset):
        key = self._buf[offset:]
        return key[: key.find(0)]

    def _Align(self, alignment):
        byte_width = 1 << alignment
        self._buf.extend(b"\x00" * _PaddingBytes(len(self._buf), byte_width))
        return byte_width

    def _Write(self, fmt, value, byte_width):
        self._buf.extend(_Pack(fmt, value, byte_width))

    def _WriteVector(self, fmt, values, byte_width):
        self._buf.extend(_PackVector(fmt, values, byte_width))

    def _WriteOffset(self, offset, byte_width):
        relative_offset = len(self._buf) - offset
        assert byte_width == 8 or relative_offset < (1 << (8 * byte_width))
        self._Write(U, relative_offset, byte_width)

    def _WriteAny(self, value, byte_width):
        fmt = {
            Type.NULL: U,
            Type.BOOL: U,
            Type.INT: I,
            Type.UINT: U,
            Type.FLOAT: F,
        }.get(value.Type)
        if fmt:
            self._Write(fmt, value.Value, byte_width)
        else:
            self._WriteOffset(value.Value, byte_width)

    def _WriteBlob(self, data, append_zero, type_):
        bit_width = BitWidth.U(len(data))
        byte_width = self._Align(bit_width)
        self._Write(U, len(data), byte_width)
        loc = len(self._buf)
        self._buf.extend(data)
        if append_zero:
            self._buf.append(0)
        self._stack.append(Value(loc, type_, bit_width))
        return loc

    def _WriteScalarVector(self, element_type, byte_width, elements, fixed):
        """Writes scalar vector elements to the underlying buffer."""
        bit_width = BitWidth.B(byte_width)
        # If you get this exception, you're trying to write a vector with a size
        # field that is bigger than the scalars you're trying to write (e.g. a
        # byte vector > 255 elements). For such types, write a "blob" instead.
        if BitWidth.U(len(elements)) > bit_width:
            raise ValueError("too many elements for the given byte_width")

        self._Align(bit_width)
        if not fixed:
            self._Write(U, len(elements), byte_width)

        loc = len(self._buf)

        fmt = {Type.INT: I, Type.UINT: U, Type.FLOAT: F}.get(element_type)
        if not fmt:
            raise TypeError("unsupported element_type")
        self._WriteVector(fmt, elements, byte_width)

        type_ = Type.ToTypedVector(element_type, len(elements) if fixed else 0)
        self._stack.append(Value(loc, type_, bit_width))
        return loc

    def _CreateVector(self, elements, typed, fixed, keys=None):
        """Writes vector elements to the underlying buffer."""
        length = len(elements)

        if fixed and not typed:
            raise ValueError("fixed vector must be typed")

        # Figure out smallest bit width we can store this vector with.
        bit_width = max(self._force_min_bit_width, BitWidth.U(length))
        prefix_elems = 1  # Vector size
        if keys:
            bit_width = max(bit_width, keys.ElemWidth(len(self._buf)))
            prefix_elems += 2  # Offset to the keys vector and its byte width.

        vector_type = Type.KEY
        # Check bit widths and types for all elements.
        for i, e in enumerate(elements):
            bit_width = max(bit_width, e.ElemWidth(len(self._buf), prefix_elems + i))

            if typed:
                if i == 0:
                    vector_type = e.Type
                else:
                    if vector_type != e.Type:
                        raise RuntimeError(
                            "typed vector elements must be of the same type"
                        )

        if fixed and not Type.IsFixedTypedVectorElementType(vector_type):
            raise RuntimeError("must be fixed typed vector element type")

        byte_width = self._Align(bit_width)
        # Write vector. First the keys width/offset if available, and size.
        if keys:
            self._WriteOffset(keys.Value, byte_width)
            self._Write(U, 1 << keys.MinBitWidth, byte_width)

        if not fixed:
            self._Write(U, length, byte_width)

        # Then the actual data.
        loc = len(self._buf)
        for e in elements:
            self._WriteAny(e, byte_width)

        # Then the types.
        if not typed:
            for e in elements:
                self._buf.append(e.StoredPackedType(bit_width))

        if keys:
            type_ = Type.MAP
        else:
            if typed:
                type_ = Type.ToTypedVector(vector_type, length if fixed else 0)
            else:
                type_ = Type.VECTOR

        return Value(loc, type_, bit_width)

    def _PushIndirect(self, value, type_, bit_width):
        byte_width = self._Align(bit_width)
        loc = len(self._buf)
        fmt = {Type.INDIRECT_INT: I, Type.INDIRECT_UINT: U, Type.INDIRECT_FLOAT: F}[
            type_
        ]
        self._Write(fmt, value, byte_width)
        self._stack.append(Value(loc, type_, bit_width))

    @InMapForString
    def String(self, value):
        """Encodes string value."""
        reset_to = len(self._buf)
        encoded = value.encode("utf-8")
        loc = self._WriteBlob(encoded, append_zero=True, type_=Type.STRING)
        if self._share_strings:
            prev_loc = self._string_pool.FindOrInsert(encoded, loc)
            if prev_loc is not None:
                del self._buf[reset_to:]
                self._stack[-1]._value = loc = prev_loc  # pylint: disable=protected-access

        return loc

    @InMap
    def Blob(self, value):
        """Encodes binary blob value.

        Args:
          value: A byte/bytearray value to encode

        Returns:
          Offset of the encoded value in underlying the byte buffer.
        """
        return self._WriteBlob(value, append_zero=False, type_=Type.BLOB)

    def Key(self, value):
        """Encodes key value.

        Args:
          value: A byte/bytearray/str value to encode. Byte object must not contain
            zero bytes. String object must be convertible to ASCII.

        Returns:
          Offset of the encoded value in the underlying byte buffer.
        """
        if isinstance(value, (bytes, bytearray)):
            encoded = value
        else:
            encoded = value.encode("ascii")

        if 0 in encoded:
            raise ValueError("key contains zero byte")

        loc = len(self._buf)
        self._buf.extend(encoded)
        self._buf.append(0)
        if self._share_keys:
            prev_loc = self._key_pool.FindOrInsert(encoded, loc)
            if prev_loc is not None:
                del self._buf[loc:]
                loc = prev_loc

        self._stack.append(Value.Key(loc))
        return loc

    def Null(self, key=None):
        """Encodes None value."""
        if key:
            self.Key(key)
        self._stack.append(Value.Null())

    @InMap
    def Bool(self, value):
        """Encodes boolean value.

        Args:
          value: A boolean value.
        """
        self._stack.append(Value.Bool(value))

    @InMap
    def Int(self, value, byte_width=0):
        """Encodes signed integer value.

        Args:
          value: A signed integer value.
          byte_width: Number of bytes to use: 1, 2, 4, or 8.
        """
        bit_width = BitWidth.I(value) if byte_width == 0 else BitWidth.B(byte_width)
        self._stack.append(Value.Int(value, bit_width))

    @InMap
    def IndirectInt(self, value, byte_width=0):
        """Encodes signed integer value indirectly.

        Args:
          value: A signed integer value.
          byte_width: Number of bytes to use: 1, 2, 4, or 8.
        """
        bit_width = BitWidth.I(value) if byte_width == 0 else BitWidth.B(byte_width)
        self._PushIndirect(value, Type.INDIRECT_INT, bit_width)

    @InMap
    def UInt(self, value, byte_width=0):
        """Encodes unsigned integer value.

        Args:
          value: An unsigned integer value.
          byte_width: Number of bytes to use: 1, 2, 4, or 8.
        """
        bit_width = BitWidth.U(value) if byte_width == 0 else BitWidth.B(byte_width)
        self._stack.append(Value.UInt(value, bit_width))

    @InMap
    def IndirectUInt(self, value, byte_width=0):
        """Encodes unsigned integer value indirectly.

        Args:
          value: An unsigned integer value.
          byte_width: Number of bytes to use: 1, 2, 4, or 8.
        """
        bit_width = BitWidth.U(value) if byte_width == 0 else BitWidth.B(byte_width)
        self._PushIndirect(value, Type.INDIRECT_UINT, bit_width)

    @InMap
    def Float(self, value, byte_width=0):
        """Encodes floating point value.

        Args:
          value: A floating point value.
          byte_width: Number of bytes to use: 4 or 8.
        """
        bit_width = BitWidth.F(value) if byte_width == 0 else BitWidth.B(byte_width)
        self._stack.append(Value.Float(value, bit_width))

    @InMap
    def IndirectFloat(self, value, byte_width=0):
        """Encodes floating point value indirectly.

        Args:
          value: A floating point value.
          byte_width: Number of bytes to use: 4 or 8.
        """
        bit_width = BitWidth.F(value) if byte_width == 0 else BitWidth.B(byte_width)
        self._PushIndirect(value, Type.INDIRECT_FLOAT, bit_width)

    def _StartVector(self):
        """Starts vector construction."""
        return len(self._stack)

    def _EndVector(self, start, typed, fixed):
        """Finishes vector construction by encodung its elements."""
        vec = self._CreateVector(self._stack[start:], typed, fixed)
        del self._stack[start:]
        self._stack.append(vec)
        return vec.Value

    @contextlib.contextmanager
    def Vector(self, key=None):
        if key:
            self.Key(key)

        try:
            start = self._StartVector()
            yield self
        finally:
            self._EndVector(start, typed=False, fixed=False)

    @InMap
    def VectorFromElements(self, elements):
        """Encodes sequence of any elements as a vector.

        Args:
          elements: sequence of elements, they may have different types.
        """
        with self.Vector():
            for e in elements:
                self.Add(e)

    @contextlib.contextmanager
    def TypedVector(self, key=None):
        if key:
            self.Key(key)

        try:
            start = self._StartVector()
            yield self
        finally:
            self._EndVector(start, typed=True, fixed=False)

    @InMap
    def TypedVectorFromElements(self, elements, element_type=None):
        """Encodes sequence of elements of the same type as typed vector.

        Args:
          elements: Sequence of elements, they must be of the same type.
          element_type: Suggested element type. Setting it to None means determining
            correct value automatically based on the given elements.
        """
        if isinstance(elements, array.array):
            if elements.typecode == "f":
                self._WriteScalarVector(Type.FLOAT, 4, elements, fixed=False)
            elif elements.typecode == "d":
                self._WriteScalarVector(Type.FLOAT, 8, elements, fixed=False)
            elif elements.typecode in ("b", "h", "i", "l", "q"):
                self._WriteScalarVector(
                    Type.INT, elements.itemsize, elements, fixed=False
                )
            elif elements.typecode in ("B", "H", "I", "L", "Q"):
                self._WriteScalarVector(
                    Type.UINT, elements.itemsize, elements, fixed=False
                )
            else:
                raise ValueError("unsupported array typecode: %s" % elements.typecode)
        else:
            add = self.Add if element_type is None else self.Adder(element_type)
            with self.TypedVector():
                for e in elements:
                    add(e)

    @InMap
    def FixedTypedVectorFromElements(self, elements, element_type=None, byte_width=0):
        """Encodes sequence of elements of the same type as fixed typed vector.

        Args:
          elements: Sequence of elements, they must be of the same type. Allowed
            types are `Type.INT`, `Type.UINT`, `Type.FLOAT`. Allowed number of
            elements are 2, 3, or 4.
          element_type: Suggested element type. Setting it to None means determining
            correct value automatically based on the given elements.
          byte_width: Number of bytes to use per element. For `Type.INT` and
            `Type.UINT`: 1, 2, 4, or 8. For `Type.FLOAT`: 4 or 8. Setting it to 0
            means determining correct value automatically based on the given
            elements.
        """
        if not 2 <= len(elements) <= 4:
            raise ValueError("only 2, 3, or 4 elements are supported")

        types = {type(e) for e in elements}
        if len(types) != 1:
            raise TypeError("all elements must be of the same type")

        (type_,) = types

        if element_type is None:
            element_type = {int: Type.INT, float: Type.FLOAT}.get(type_)
            if not element_type:
                raise TypeError("unsupported element_type: %s" % type_)

        if byte_width == 0:
            width = {
                Type.UINT: BitWidth.U,
                Type.INT: BitWidth.I,
                Type.FLOAT: BitWidth.F,
            }[element_type]
            byte_width = 1 << max(width(e) for e in elements)

        self._WriteScalarVector(element_type, byte_width, elements, fixed=True)

    def _StartMap(self):
        """Starts map construction."""
        return len(self._stack)

    def _EndMap(self, start):
        """Finishes map construction by encodung its elements."""
        # Interleaved keys and values on the stack.
        stack = self._stack[start:]

        if len(stack) % 2 != 0:
            raise RuntimeError("must be even number of keys and values")

        for key in stack[::2]:
            if key.Type is not Type.KEY:
                raise RuntimeError("all map keys must be of %s type" % Type.KEY)

        pairs = zip(stack[::2], stack[1::2])  # [(key, value), ...]
        pairs = sorted(pairs, key=lambda pair: self._ReadKey(pair[0].Value))

        del self._stack[start:]
        for pair in pairs:
            self._stack.extend(pair)

        keys = self._CreateVector(self._stack[start::2], typed=True, fixed=False)
        values = self._CreateVector(
            self._stack[start + 1 :: 2], typed=False, fixed=False, keys=keys
        )

        del self._stack[start:]
        self._stack.append(values)
        return values.Value

    @contextlib.contextmanager
    def Map(self, key=None):
        if key:
            self.Key(key)

        try:
            start = self._StartMap()
            yield self
        finally:
            self._EndMap(start)

    def MapFromElements(self, elements):
        start = self._StartMap()
        for k, v in elements.items():
            self.Key(k)
            self.Add(v)
        self._EndMap(start)

    def Adder(self, type_):
        return {
            Type.BOOL: self.Bool,
            Type.INT: self.Int,
            Type.INDIRECT_INT: self.IndirectInt,
            Type.UINT: self.UInt,
            Type.INDIRECT_UINT: self.IndirectUInt,
            Type.FLOAT: self.Float,
            Type.INDIRECT_FLOAT: self.IndirectFloat,
            Type.KEY: self.Key,
            Type.BLOB: self.Blob,
            Type.STRING: self.String,
        }[type_]

    @InMapForString
    def Add(self, value):
        """Encodes value of any supported type."""
        if value is None:
            self.Null()
        elif isinstance(value, bool):
            self.Bool(value)
        elif isinstance(value, int):
            self.Int(value)
        elif isinstance(value, float):
            self.Float(value)
        elif isinstance(value, str):
            self.String(value)
        elif isinstance(value, (bytes, bytearray)):
            self.Blob(value)
        elif isinstance(value, dict):
            with self.Map():
                for k, v in value.items():
                    self.Key(k)
                    self.Add(v)
        elif isinstance(value, array.array):
            self.TypedVectorFromElements(value)
        elif _IsIterable(value):
            self.VectorFromElements(value)
        else:
            raise TypeError("unsupported python type: %s" % type(value))

    @property
    def LastValue(self):
        return self._stack[-1]

    @InMap
    def ReuseValue(self, value):
        self._stack.append(value)


def GetRoot(buf):
    """Returns root `Ref` object for the given buffer."""
    if len(buf) < 3:
        raise ValueError("buffer is too small")
    byte_width = buf[-1]
    return Ref.PackedType(Buf(buf, -(2 + byte_width)), byte_width, packed_type=buf[-2])


def Dumps(obj):
    """Returns bytearray with the encoded python object."""
    fbb = Builder()
    fbb.Add(obj)
    return fbb.Finish()


def Loads(buf):
    """Returns python object decoded from the buffer."""
    return GetRoot(buf).Value
