"""
Convert between bytestreams and higher-level AMQP types.

2007-11-05 Barry Pederson <bp@barryp.org>

"""
# Copyright (C) 2007 Barry Pederson <bp@barryp.org>
#
# This library is free software; you can redistribute it and/or
# modify it under the terms of the GNU Lesser General Public
# License as published by the Free Software Foundation; either
# version 2.1 of the License, or (at your option) any later version.
#
# This library is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
# Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public
# License along with this library; if not, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301
from __future__ import absolute_import

import sys

from datetime import datetime
from decimal import Decimal
from struct import pack, unpack
from time import mktime

IS_PY3K = sys.version_info[0] >= 3

if IS_PY3K:
    def byte(n):
        return bytes([n])
else:
    byte = chr

try:
    from io import BytesIO
except ImportError:  # Py2.5
    try:
        from cStringIO import StringIO as BytesIO  # noqa
    except ImportError:
        from StringIO import StringIO as BytesIO   # noqa

try:
    bytes
except NameError:
    # Python 2.5 and lower
    bytes = str


class AMQPReader(object):
    """Read higher-level AMQP types from a bytestream."""
    def __init__(self, source):
        """Source should be either a file-like object with a read() method, or
        a plain (non-unicode) string."""
        if isinstance(source, bytes):
            self.input = BytesIO(source)
        elif hasattr(source, 'read'):
            self.input = source
        else:
            raise ValueError(
                'AMQPReader needs a file-like object or plain string')

        self.bitcount = self.bits = 0

    def close(self):
        self.input.close()

    def read(self, n):
        """Read n bytes."""
        self.bitcount = self.bits = 0
        return self.input.read(n)

    def read_bit(self):
        """Read a single boolean value."""
        if not self.bitcount:
            self.bits = ord(self.input.read(1))
            self.bitcount = 8
        result = (self.bits & 1) == 1
        self.bits >>= 1
        self.bitcount -= 1
        return result

    def read_octet(self):
        """Read one byte, return as an integer"""
        self.bitcount = self.bits = 0
        return unpack('B', self.input.read(1))[0]

    def read_short(self):
        """Read an unsigned 16-bit integer"""
        self.bitcount = self.bits = 0
        return unpack('>H', self.input.read(2))[0]

    def read_long(self):
        """Read an unsigned 32-bit integer"""
        self.bitcount = self.bits = 0
        return unpack('>I', self.input.read(4))[0]

    def read_longlong(self):
        """Read an unsigned 64-bit integer"""
        self.bitcount = self.bits = 0
        return unpack('>Q', self.input.read(8))[0]

    def read_float(self):
        """Read float value."""
        self.bitcount = self.bits = 0
        return unpack('>d', self.input.read(8))[0]

    def read_shortstr(self):
        """Read a short string that's stored in up to 255 bytes.

        The encoding isn't specified in the AMQP spec, so
        assume it's utf-8

        """
        self.bitcount = self.bits = 0
        slen = unpack('B', self.input.read(1))[0]
        return self.input.read(slen).decode('utf-8')

    def read_longstr(self):
        """Read a string that's up to 2**32 bytes.

        The encoding isn't specified in the AMQP spec, so
        assume it's utf-8

        """
        self.bitcount = self.bits = 0
        slen = unpack('>I', self.input.read(4))[0]
        return self.input.read(slen).decode('utf-8')

    def read_table(self):
        """Read an AMQP table, and return as a Python dictionary."""
        self.bitcount = self.bits = 0
        tlen = unpack('>I', self.input.read(4))[0]
        table_data = AMQPReader(self.input.read(tlen))
        result = {}
        while table_data.input.tell() < tlen:
            name = table_data.read_shortstr()
            ftype = ord(table_data.input.read(1))
            if ftype == 83:  # 'S'
                val = table_data.read_longstr()
            elif ftype == 73:  # 'I'
                val = unpack('>i', table_data.input.read(4))[0]
            elif ftype == 68:  # 'D'
                d = table_data.read_octet()
                n = unpack('>i', table_data.input.read(4))[0]
                val = Decimal(n) / Decimal(10 ** d)
            elif ftype == 84:  # 'T'
                val = table_data.read_timestamp()
            elif ftype == 70:  # 'F'
                val = table_data.read_table()  # recurse
            elif ftype == 116:
                val = table_data.read_bit()
            elif ftype == 100:
                val = table_data.read_float()
            else:
                raise ValueError('Unknown table item type: %s' % repr(ftype))
            result[name] = val
        return result

    def read_timestamp(self):
        """Read and AMQP timestamp, which is a 64-bit integer representing
        seconds since the Unix epoch in 1-second resolution.

        Return as a Python datetime.datetime object,
        expressed as localtime.

        """
        return datetime.fromtimestamp(self.read_longlong())


class AMQPWriter(object):
    """Convert higher-level AMQP types to bytestreams."""

    def __init__(self, dest=None):
        """dest may be a file-type object (with a write() method).  If None
        then a BytesIO is created, and the contents can be accessed with
        this class's getvalue() method."""
        self.out = BytesIO() if dest is None else dest
        self.bits = []
        self.bitcount = 0

    def _flushbits(self):
        if self.bits:
            out = self.out
            for b in self.bits:
                out.write(pack('B', b))
            self.bits = []
            self.bitcount = 0

    def close(self):
        """Pass through if possible to any file-like destinations."""
        try:
            self.out.close()
        except AttributeError:
            pass

    def flush(self):
        """Pass through if possible to any file-like destinations."""
        try:
            self.out.flush()
        except AttributeError:
            pass

    def getvalue(self):
        """Get what's been encoded so far if we're working with a BytesIO."""
        self._flushbits()
        return self.out.getvalue()

    def write(self, s):
        """Write a plain Python string with no special encoding in Python 2.x,
        or bytes in Python 3.x"""
        self._flushbits()
        self.out.write(s)

    def write_bit(self, b):
        """Write a boolean value."""
        b = 1 if b else 0
        shift = self.bitcount % 8
        if shift == 0:
            self.bits.append(0)
        self.bits[-1] |= (b << shift)
        self.bitcount += 1

    def write_octet(self, n):
        """Write an integer as an unsigned 8-bit value."""
        if n < 0 or n > 255:
            raise ValueError('Octet %r out of range 0..255' % (n, ))
        self._flushbits()
        self.out.write(pack('B', n))

    def write_short(self, n):
        """Write an integer as an unsigned 16-bit value."""
        if n < 0 or n > 65535:
            raise ValueError('Octet %r out of range 0..65535' % (n, ))
        self._flushbits()
        self.out.write(pack('>H', int(n)))

    def write_long(self, n):
        """Write an integer as an unsigned2 32-bit value."""
        if n < 0 or n >= 2 ** 32:
            raise ValueError('Octet %r out of range 0..2**31-1' % (n, ))
        self._flushbits()
        self.out.write(pack('>I', n))

    def write_longlong(self, n):
        """Write an integer as an unsigned 64-bit value."""
        if n < 0 or n >= 2 ** 64:
            raise ValueError('Octet %r out of range 0..2**64-1' % (n, ))
        self._flushbits()
        self.out.write(pack('>Q', n))

    def write_shortstr(self, s):
        """Write a string up to 255 bytes long (after any encoding).

        If passed a unicode string, encode with UTF-8.

        """
        self._flushbits()
        if isinstance(s, unicode):
            s = s.encode('utf-8')
        if len(s) > 255:
            raise ValueError('String too long (%r)' % (len(s), ))
        self.write_octet(len(s))
        self.out.write(s)

    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.

        """
        self._flushbits()
        if isinstance(s, unicode):
            s = s.encode('utf-8')
        self.write_long(len(s))
        self.out.write(s)

    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."""
        self._flushbits()
        table_data = AMQPWriter()
        for k, v in d.iteritems():
            table_data.write_shortstr(k)
            if isinstance(v, basestring):
                if isinstance(v, unicode):
                    v = v.encode('utf-8')
                table_data.write(byte(83))  # 'S'
                table_data.write_longstr(v)
            elif isinstance(v, bool):
                table_data.write(pack('>cB', 't', int(v)))
            elif isinstance(v, float):
                table_data.write(pack('>cd', 'd', v))
            elif isinstance(v, (int, long)):
                table_data.write(pack('>ci', 'I', v))
            elif isinstance(v, Decimal):
                table_data.write(byte(68))  # 'D'
                sign, digits, exponent = v.as_tuple()
                v = 0
                for d in digits:
                    v = (v * 10) + d
                if sign:
                    v = -v
                table_data.write_octet(-exponent)
                table_data.write(pack('>i', v))
            elif isinstance(v, datetime):
                table_data.write(byte(84))  # 'T'
                table_data.write_timestamp(v)
                ## FIXME: timezone ?
            elif isinstance(v, dict):
                table_data.write(byte(70))  # 'F'
                table_data.write_table(v)
            else:
                raise ValueError('%r not serializable in AMQP' % (v, ))
        table_data = table_data.getvalue()
        self.write_long(len(table_data))
        self.out.write(table_data)

    def write_timestamp(self, v):
        """Write out a Python datetime.datetime object as a 64-bit integer
        representing seconds since the Unix epoch."""
        self.out.write(pack('>q', long(mktime(v.timetuple()))))


class GenericContent(object):
    """Abstract base class for AMQP content.

    Subclasses should override the PROPERTIES attribute.

    """
    PROPERTIES = [('dummy', 'shortstr')]

    def __init__(self, **props):
        """Save the properties appropriate to this AMQP content type
        in a 'properties' dictionary."""
        d = {}
        for propname, _ in self.PROPERTIES:
            if propname in props:
                d[propname] = props[propname]
            # FIXME: should we ignore unknown properties?

        self.properties = d

    def __eq__(self, other):
        """Check if this object has the same properties as another
        content object."""
        try:
            return self.properties == other.properties
        except AttributeError:
            return NotImplemented

    def __getattr__(self, name):
        """Look for additional properties in the 'properties'
        dictionary, and if present - the 'delivery_info'
        dictionary."""
        if name == '__setstate__':
            # Allows pickling/unpickling to work
            raise AttributeError('__setstate__')

        if name in self.properties:
            return self.properties[name]

        if 'delivery_info' in self.__dict__ \
                and name in self.delivery_info:
            return self.delivery_info[name]

        raise AttributeError(name)

    def _load_properties(self, raw_bytes):
        """Given the raw bytes containing the property-flags and property-list
        from a content-frame-header, parse and insert into a dictionary
        stored in this object as an attribute named 'properties'."""
        r = AMQPReader(raw_bytes)

        #
        # Read 16-bit shorts until we get one with a low bit set to zero
        #
        flags = []
        while 1:
            flag_bits = r.read_short()
            flags.append(flag_bits)
            if flag_bits & 1 == 0:
                break

        shift = 0
        d = {}
        for key, proptype in self.PROPERTIES:
            if shift == 0:
                if not flags:
                    break
                flag_bits, flags = flags[0], flags[1:]
                shift = 15
            if flag_bits & (1 << shift):
                d[key] = getattr(r, 'read_' + proptype)()
            shift -= 1

        self.properties = d

    def _serialize_properties(self):
        """serialize the 'properties' attribute (a dictionary) into
        the raw bytes making up a set of property flags and a
        property list, suitable for putting into a content frame header."""
        shift = 15
        flag_bits = 0
        flags = []
        raw_bytes = AMQPWriter()
        for key, proptype in self.PROPERTIES:
            val = self.properties.get(key, None)
            if val is not None:
                if shift == 0:
                    flags.append(flag_bits)
                    flag_bits = 0
                    shift = 15

                flag_bits |= (1 << shift)
                if proptype != 'bit':
                    getattr(raw_bytes, 'write_' + proptype)(val)

            shift -= 1

        flags.append(flag_bits)
        result = AMQPWriter()
        for flag_bits in flags:
            result.write_short(flag_bits)
        result.write(raw_bytes.getvalue())

        return result.getvalue()
