# 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('<%s' % 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('<%s' % 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 AsStringBytes(self):
    if self.IsString:
      return String(self._Indirect(), self._byte_width).Bytes
    elif self.IsKey:
      return self.AsKeyBytes
    else:
      raise self._ConvertError(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