import struct


class Primitive(object):
    """
    The most basic structure of the protocol.  Subclassed, never used directly.

    Used as a building block for the various actually-used primitives outlined
    in the Zookeeper jute file:

    https://github.com/apache/zookeeper/blob/trunk/src/zookeeper.jute
    """
    fmt = None

    def __init__(self, value):
        self.value = value

    def render(self):
        """
        Returns a two-element tuple with the ``struct`` format and list value.

        The value is wrapped in a list, as there are some primitives that deal
        with multiple values.  Any caller of `render()` should expect a list.
        """
        return self.fmt, [self.value]

    @classmethod
    def parse(cls, buff, offset):
        """
        Given a buffer and offset, returns the parsed value and new offset.

        Uses the ``format`` class attribute to unpack the data from the buffer
        and determine the used up number of bytes.
        """
        primitive_struct = struct.Struct("!" + cls.fmt)

        value = primitive_struct.unpack_from(buff, offset)[0]
        offset += primitive_struct.size

        return value, offset

    def __eq__(self, other):
        """
        Basic equality method that tests equality of the ``value`` attributes.
        """
        return self.value == other.value

    def __str__(self):
        return "%s(%s)" % (self.__class__.__name__, self.value)


class VariablePrimitive(Primitive):
    """
    Base primitive for variable-length scalar primitives (strings and bytes).
    """
    size_primitive = None

    def render_value(self, value):
        raise NotImplementedError

    @classmethod
    def parse_value(cls, value):
        raise NotImplementedError

    def render(self):
        """
        Returns the ``struct`` format and list of the size and value.

        The format is derived from the size primitive and the length of the
        resulting encoded value (e.g. the format for a string of 'foo' ends
        up as 'h3s'.

        .. note ::
          The value is expected to be string-able (wrapped in ``str()``) and is
          then encoded as UTF-8.
        """
        size_format = self.size_primitive.fmt

        if self.value is None:
            return size_format, [-1]

        value = self.render_value(self.value)

        size = len(value)

        fmt = "%s%ds" % (size_format, size)

        return fmt, [size, value]

    @classmethod
    def parse(cls, buff, offset):
        """
        Given a buffer and offset, returns the parsed value and new offset.

        Parses the ``size_primitive`` first to determine how many more bytes to
        consume to extract the value.
        """
        size, offset = cls.size_primitive.parse(buff, offset)
        if size == -1:
            return None, offset

        var_struct = struct.Struct("!%ds" % size)

        value = var_struct.unpack_from(buff, offset)[0]
        value = cls.parse_value(value)
        offset += var_struct.size

        return value, offset


class Bool(Primitive):
    """
    Represents a boolean (true or false) value.

    Renders as an unsigned char (1 byte).
    """
    fmt = "?"


class Byte(Primitive):
    """
    Represents a single 8-bit byte.
    """
    fmt = "b"


class Int(Primitive):
    """
    Represents an 32-bit signed integer.
    """
    fmt = "i"


class Long(Primitive):
    """
    Represents an 64-bit signed integer.
    """
    fmt = "q"


class Float(Primitive):
    """
    Represents a single-precision floating poing conforming to IEEE 754.
    """
    fmt = "f"


class Double(Primitive):
    """
    Represents a double-precision floating poing conforming to IEEE 754.
    """
    fmt = "d"


class UString(VariablePrimitive):
    """
    Represents a unicode string value, length denoted by a 32-bit integer.
    """
    size_primitive = Int

    def render_value(self, value):
        return bytes(str(value).encode("utf-8"))

    @classmethod
    def parse_value(cls, value):
        return value.decode("utf-8")

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


class Buffer(VariablePrimitive):
    """
    Represents a bytestring value, length denoted by a 32-bit signed integer.
    """
    size_primitive = Int

    def render_value(self, value):
        if isinstance(value, str):
            return value.encode()
        # return bytes(value)
        return value

    @classmethod
    def parse_value(cls, value):
        return value


class Vector(Primitive):
    """
    Represents an array of any arbitrary `Primitive` or ``Part``.

    Not used directly but rather by its ``of()`` classmethod to denote an
    ``Vector.of(<something>)``.
    """
    item_class = None

    @classmethod
    def of(cls, part_class):
        """
        Creates a new class with the ``item_class`` attribute properly set.
        """
        copy = type(
            "VectorOf%s" % part_class.__name__,
            cls.__bases__, dict(cls.__dict__)
        )
        copy.item_class = part_class

        return copy

    def render(self):
        """
        Creates a composite ``struct`` format and the data to render with it.

        The format and data are prefixed with a 32-bit integer denoting the
        number of elements, after which each of the items in the array value
        are ``render()``-ed and added to the format and data as well.
        """
        value = self.value
        if value is None:
            value = []

        fmt = [Int.fmt]
        data = [len(value)]

        for item_value in value:
            if issubclass(self.item_class, Primitive):
                item = self.item_class(item_value)
            else:
                item = item_value

            item_format, item_data = item.render()
            fmt.extend(item_format)
            data.extend(item_data)

        return "".join(fmt), data

    @classmethod
    def parse(cls, buff, offset):
        """
        Parses a raw buffer at offset and returns the resulting array value.

        Starts off by `parse()`-ing the 32-bit element count, followed by
        parsing items out of the buffer "count" times.
        """
        count, offset = Int.parse(buff, offset)

        values = []
        for _ in range(count):
            value, new_offset = cls.item_class.parse(buff, offset)

            values.append(value)
            offset = new_offset

        return values, offset

    def __str__(self):
        return "%s[%s]" % (
            self.item_class.__name__, ", ".join(map(str, self.value))
        )
