'''
Copyright (c) 2011-2014, Agora Games, LLC All rights reserved.

https://github.com/agoragames/haigha/blob/master/LICENSE.txt
'''

from struct import Struct
from calendar import timegm
from datetime import datetime
from decimal import Decimal
from operator import xor


class Writer(object):

    """
    Implements writing of structured AMQP data. Buffers data directly to a
    bytearray or a buffer supplied in the constructor. The buffer must
    supply append, extend and struct.pack_into semantics.
    """

    def __init__(self, buf=None):
        if buf is not None:
            self._output_buffer = buf
        else:
            self._output_buffer = bytearray()

    def __str__(self):
        return ''.join([
            '\\x%s' % (chr(c).encode('hex')) for c in self._output_buffer])

    __repr__ = __str__

    def __eq__(self, other):
        if isinstance(other, Writer):
            return self._output_buffer == other._output_buffer
        return False

    def buffer(self):
        '''
        Get the buffer that this has written to. Returns bytearray.
        '''
        return self._output_buffer

    def write(self, s):
        """
        Write a plain Python string, with no special encoding.
        """
        self._output_buffer.extend(s)
        return self

    def write_bits(self, *args):
        '''
        Write multiple bits in a single byte field. The bits will be written in
        little-endian order, but should be supplied in big endian order. Will
        raise ValueError when more than 8 arguments are supplied.

        write_bits(True, False) => 0x02
        '''
        # Would be nice to make this a bit smarter
        if len(args) > 8:
            raise ValueError("Can only write 8 bits at a time")

        self._output_buffer.append(chr(
            reduce(lambda x, y: xor(x, args[y] << y), xrange(len(args)), 0)))

        return self

    def write_bit(self, b, pack=Struct('B').pack):
        '''
        Write a single bit. Convenience method for single bit args.
        '''
        self._output_buffer.append(pack(True if b else False))
        return self

    def write_octet(self, n, pack=Struct('B').pack):
        """
        Write an integer as an unsigned 8-bit value.
        """
        if 0 <= n <= 255:
            self._output_buffer.append(pack(n))
        else:
            raise ValueError('Octet %d out of range 0..255', n)
        return self

    def write_short(self, n, pack=Struct('>H').pack):
        """
        Write an integer as an unsigned 16-bit value.
        """
        if 0 <= n <= 0xFFFF:
            self._output_buffer.extend(pack(n))
        else:
            raise ValueError('Short %d out of range 0..0xFFFF', n)
        return self

    def write_short_at(self, n, pos, pack_into=Struct('>H').pack_into):
        '''
        Write an unsigned 16bit value at a specific position in the buffer.
        Used for writing tables and frames.
        '''
        if 0 <= n <= 0xFFFF:
            pack_into(self._output_buffer, pos, n)
        else:
            raise ValueError('Short %d out of range 0..0xFFFF', n)
        return self

    def write_long(self, n, pack=Struct('>I').pack):
        """
        Write an integer as an unsigned 32-bit value.
        """
        if 0 <= n <= 0xFFFFFFFF:
            self._output_buffer.extend(pack(n))
        else:
            raise ValueError('Long %d out of range 0..0xFFFFFFFF', n)
        return self

    def write_long_at(self, n, pos, pack_into=Struct('>I').pack_into):
        '''
        Write an unsigned 32bit value at a specific position in the buffer.
        Used for writing tables and frames.
        '''
        if 0 <= n <= 0xFFFFFFFF:
            pack_into(self._output_buffer, pos, n)
        else:
            raise ValueError('Long %d out of range 0..0xFFFFFFFF', n)
        return self

    def write_longlong(self, n, pack=Struct('>Q').pack):
        """
        Write an integer as an unsigned 64-bit value.
        """
        if 0 <= n <= 0xFFFFFFFFFFFFFFFF:
            self._output_buffer.extend(pack(n))
        else:
            raise ValueError(
                'Longlong %d out of range 0..0xFFFFFFFFFFFFFFFF', n)
        return self

    def write_shortstr(self, s):
        """
        Write a string up to 255 bytes long after encoding.  If passed
        a unicode string, encode as UTF-8.
        """
        if isinstance(s, unicode):
            s = s.encode('utf-8')
        self.write_octet(len(s))
        self.write(s)
        return self

    def write_longstr(self, s):
        """
        Write a string up to 2**32 bytes long after encoding.  If passed
        a unicode string, encode as UTF-8.
        """
        if isinstance(s, unicode):
            s = s.encode('utf-8')
        self.write_long(len(s))
        self.write(s)
        return self

    def write_timestamp(self, t, pack=Struct('>Q').pack):
        """
        Write out a Python datetime.datetime object as a 64-bit integer
        representing seconds since the Unix UTC epoch.
        """
        # Double check timestamp, can't imagine why it would be signed
        self._output_buffer.extend(pack(long(timegm(t.timetuple()))))
        return self

    # NOTE: coding to http://dev.rabbitmq.com/wiki/Amqp091Errata#section_3 and
    # NOT spec 0.9.1. It seems that Rabbit and other brokers disagree on this
    # section for now.
    def write_table(self, d):
        """
        Write out a Python dictionary made of up string keys, and values
        that are strings, signed integers, Decimal, datetime.datetime, or
        sub-dictionaries following the same constraints.
        """
        # HACK: encoding of AMQP tables is broken because it requires the
        # length of the /encoded/ data instead of the number of items. To
        # support streaming, fiddle with cursor position, rewinding to write
        # the real length of the data. Generally speaking, I'm not a fan of
        # the AMQP encoding scheme, it could be much faster.
        table_len_pos = len(self._output_buffer)
        self.write_long(0)
        table_data_pos = len(self._output_buffer)

        for key, value in d.iteritems():
            self._write_item(key, value)

        table_end_pos = len(self._output_buffer)
        table_len = table_end_pos - table_data_pos

        self.write_long_at(table_len, table_len_pos)
        return self

    def _write_item(self, key, value):
        self.write_shortstr(key)
        self._write_field(value)

    def _write_field(self, value):
        writer = self.field_type_map.get(type(value))
        if writer:
            writer(self, value)
        else:
            for kls, writer in self.field_type_map.items():
                if isinstance(value, kls):
                    writer(self, value)
                    break
            else:
                # Write a None because we've already written a key
                self._field_none(value)

    def _field_bool(self, val, pack=Struct('B').pack):
        self._output_buffer.append('t')
        self._output_buffer.append(pack(True if val else False))

    def _field_int(self, val, short_pack=Struct('>h').pack,
                   int_pack=Struct('>i').pack, long_pack=Struct('>q').pack):
        if -2 ** 15 <= val < 2 ** 15:
            self._output_buffer.append('s')
            self._output_buffer.extend(short_pack(val))
        elif -2 ** 31 <= val < 2 ** 31:
            self._output_buffer.append('I')
            self._output_buffer.extend(int_pack(val))
        else:
            self._output_buffer.append('l')
            self._output_buffer.extend(long_pack(val))

    def _field_double(self, val, pack=Struct('>d').pack):
        self._output_buffer.append('d')
        self._output_buffer.extend(pack(val))

    # Coding to http://dev.rabbitmq.com/wiki/Amqp091Errata#section_3 which
    # differs from spec in that the value is signed.
    def _field_decimal(self, val, exp_pack=Struct('B').pack,
                       dig_pack=Struct('>i').pack):
        self._output_buffer.append('D')
        sign, digits, exponent = val.as_tuple()
        v = 0
        for d in digits:
            v = (v * 10) + d
        if sign:
            v = -v
        self._output_buffer.append(exp_pack(-exponent))
        self._output_buffer.extend(dig_pack(v))

    def _field_str(self, val):
        self._output_buffer.append('S')
        self.write_longstr(val)

    def _field_unicode(self, val):
        val = val.encode('utf-8')
        self._output_buffer.append('S')
        self.write_longstr(val)

    def _field_timestamp(self, val):
        self._output_buffer.append('T')
        self.write_timestamp(val)

    def _field_table(self, val):
        self._output_buffer.append('F')
        self.write_table(val)

    def _field_none(self, val):
        self._output_buffer.append('V')

    def _field_bytearray(self, val):
        self._output_buffer.append('x')
        self.write_longstr(val)

    def _field_iterable(self, val):
        self._output_buffer.append('A')
        for x in val:
            self._write_field(x)

    field_type_map = {
        bool: _field_bool,
        int: _field_int,
        long: _field_int,
        float: _field_double,
        Decimal: _field_decimal,
        str: _field_str,
        unicode: _field_unicode,
        datetime: _field_timestamp,
        dict: _field_table,
        type(None): _field_none,
        bytearray: _field_bytearray,
    }

    # 0.9.1 spec mapping
    # field_type_map = {
    #   bool      : _field_bool,
    #   int       : _field_int,
    #   long      : _field_int,
    #   float     : _field_double,
    #   Decimal   : _field_decimal,
    #   str       : _field_str,
    #   unicode   : _field_unicode,
    #   datetime  : _field_timestamp,
    #   dict      : _field_table,
    #   None      : _field_none,
    #   bytearray : _field_bytearray,
    #   list      : _field_iterable,
    #   tuple     : _field_iterable,
    #   set       : _field_iterable,
    # }
