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

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

from collections import deque

from haigha.message import Message
from haigha.writer import Writer
from haigha.frames.method_frame import MethodFrame
from haigha.frames.header_frame import HeaderFrame
from haigha.frames.content_frame import ContentFrame
from haigha.classes.protocol_class import ProtocolClass


class BasicClass(ProtocolClass):

    '''
    Implements the AMQP Basic class
    '''

    def __init__(self, *args, **kwargs):
        super(BasicClass, self).__init__(*args, **kwargs)
        self.dispatch_map = {
            11: self._recv_qos_ok,
            21: self._recv_consume_ok,
            31: self._recv_cancel_ok,
            50: self._recv_return,
            60: self._recv_deliver,
            71: self._recv_get_response,   # see impl
            72: self._recv_get_response,   # see impl
            111: self._recv_recover_ok,
        }

        self._consumer_tag_id = 0
        self._pending_consumers = deque()
        self._consumer_cb = {}
        self._get_cb = deque()
        self._recover_cb = deque()
        self._cancel_cb = deque()

    @property
    def name(self):
        return 'basic'

    def _cleanup(self):
        '''
        Cleanup all the local data.
        '''
        self._pending_consumers = None
        self._consumer_cb = None
        self._get_cb = None
        self._recover_cb = None
        self._cancel_cb = None
        super(BasicClass, self)._cleanup()

    def _generate_consumer_tag(self):
        '''
        Generate the next consumer tag.

        The consumer tag is local to a channel, so two clients can use the
        same consumer tags.
        '''
        self._consumer_tag_id += 1
        return "channel-%d-%d" % (self.channel_id, self._consumer_tag_id)

    def qos(self, prefetch_size=0, prefetch_count=0, is_global=False):
        '''
        Set QoS on this channel.
        '''
        args = Writer()
        args.write_long(prefetch_size).\
            write_short(prefetch_count).\
            write_bit(is_global)
        self.send_frame(MethodFrame(self.channel_id, 60, 10, args))

        self.channel.add_synchronous_cb(self._recv_qos_ok)

    def _recv_qos_ok(self, _method_frame):
        # No arguments, nothing to do
        pass

    def consume(self, queue, consumer, consumer_tag='', no_local=False,
                no_ack=True, exclusive=False, nowait=True, ticket=None,
                cb=None):
        '''
        Start a queue consumer. If `cb` is supplied, will be called when
        broker confirms that consumer is registered.
        '''
        nowait = nowait and self.allow_nowait() and not cb

        if nowait and consumer_tag == '':
            consumer_tag = self._generate_consumer_tag()

        args = Writer()
        args.write_short(ticket or self.default_ticket).\
            write_shortstr(queue).\
            write_shortstr(consumer_tag).\
            write_bits(no_local, no_ack, exclusive, nowait).\
            write_table({})  # unused according to spec
        self.send_frame(MethodFrame(self.channel_id, 60, 20, args))

        if not nowait:
            self._pending_consumers.append((consumer, cb))
            self.channel.add_synchronous_cb(self._recv_consume_ok)
        else:
            self._consumer_cb[consumer_tag] = consumer

    def _recv_consume_ok(self, method_frame):
        consumer_tag = method_frame.args.read_shortstr()
        consumer, cb = self._pending_consumers.popleft()

        self._consumer_cb[consumer_tag] = consumer
        if cb:
            cb()

    def cancel(self, consumer_tag='', nowait=True, consumer=None, cb=None):
        '''
        Cancel a consumer. Can choose to delete based on a consumer tag or
        the function which is consuming.  If deleting by function, take care
        to only use a consumer once per channel.
        '''
        if consumer:
            for (tag, func) in self._consumer_cb.iteritems():
                if func == consumer:
                    consumer_tag = tag
                    break

        nowait = nowait and self.allow_nowait() and not cb

        args = Writer()
        args.write_shortstr(consumer_tag).\
            write_bit(nowait)
        self.send_frame(MethodFrame(self.channel_id, 60, 30, args))

        if not nowait:
            self._cancel_cb.append(cb)
            self.channel.add_synchronous_cb(self._recv_cancel_ok)
        else:
            try:
                del self._consumer_cb[consumer_tag]
            except KeyError:
                self.logger.warning(
                    'no callback registered for consumer tag " %s "',
                    consumer_tag)

    def _recv_cancel_ok(self, method_frame):
        consumer_tag = method_frame.args.read_shortstr()
        try:
            del self._consumer_cb[consumer_tag]
        except KeyError:
            self.logger.warning(
                'no callback registered for consumer tag " %s "', consumer_tag)

        cb = self._cancel_cb.popleft()
        if cb:
            cb()

    def publish(self, msg, exchange, routing_key, mandatory=False,
                immediate=False, ticket=None):
        '''
        publish a message.
        '''
        args = Writer()
        args.write_short(ticket or self.default_ticket).\
            write_shortstr(exchange).\
            write_shortstr(routing_key).\
            write_bits(mandatory, immediate)

        self.send_frame(MethodFrame(self.channel_id, 60, 40, args))
        self.send_frame(
            HeaderFrame(self.channel_id, 60, 0, len(msg), msg.properties))

        f_max = self.channel.connection.frame_max
        for f in ContentFrame.create_frames(self.channel_id, msg.body, f_max):
            self.send_frame(f)

    def return_msg(self, reply_code, reply_text, exchange, routing_key):
        '''
        Return a failed message.  Not named "return" because python interpreter
        can't deal with that.
        '''
        args = Writer()
        args.write_short(reply_code).\
            write_shortstr(reply_text).\
            write_shortstr(exchange).\
            write_shortstr(routing_key)

        self.send_frame(MethodFrame(self.channel_id, 60, 50, args))

    def _recv_return(self, _method_frame):
        # This seems like the right place to callback that the operation has
        # completed.
        pass

    def _recv_deliver(self, method_frame):
        msg = self._read_msg(method_frame,
                             with_consumer_tag=True, with_message_count=False)

        func = self._consumer_cb.get(msg.delivery_info['consumer_tag'], None)
        if func:
            func(msg)

    def get(self, queue, consumer=None, no_ack=True, ticket=None):
        '''
        Ask to fetch a single message from a queue.  If a consumer is supplied,
        the consumer will be called with either a Message argument, or None if
        there is no message in queue. If a synchronous transport, Message or
        None is returned.
        '''
        args = Writer()
        args.write_short(ticket or self.default_ticket).\
            write_shortstr(queue).\
            write_bit(no_ack)

        self._get_cb.append(consumer)
        self.send_frame(MethodFrame(self.channel_id, 60, 70, args))
        return self.channel.add_synchronous_cb(self._recv_get_response)

    def _recv_get_response(self, method_frame):
        '''
        Handle either get_ok or get_empty. This is a hack because the
        synchronous callback stack is expecting one method to satisfy the
        expectation. To keep that loop as tight as possible, work within
        those constraints. Use of get is not recommended anyway.
        '''
        if method_frame.method_id == 71:
            return self._recv_get_ok(method_frame)
        elif method_frame.method_id == 72:
            return self._recv_get_empty(method_frame)

    def _recv_get_ok(self, method_frame):
        msg = self._read_msg(method_frame,
                             with_consumer_tag=False, with_message_count=True)
        cb = self._get_cb.popleft()
        if cb:
            cb(msg)
        return msg

    def _recv_get_empty(self, _method_frame):
        # On empty, call back with None as the argument so that user code knows
        # it's empty and can take next action
        cb = self._get_cb.popleft()
        if cb:
            cb(None)

    def ack(self, delivery_tag, multiple=False):
        '''
        Acknowledge delivery of a message.  If multiple=True, acknowledge up-to
        and including delivery_tag.
        '''
        args = Writer()
        args.write_longlong(delivery_tag).\
            write_bit(multiple)

        self.send_frame(MethodFrame(self.channel_id, 60, 80, args))

    def reject(self, delivery_tag, requeue=False):
        '''
        Reject a message.
        '''
        args = Writer()
        args.write_longlong(delivery_tag).\
            write_bit(requeue)

        self.send_frame(MethodFrame(self.channel_id, 60, 90, args))

    def recover_async(self, requeue=False):
        '''
        Redeliver all unacknowledged messages on this channel.

        This method is deprecated in favour of the synchronous
        recover/recover-ok
        '''
        args = Writer()
        args.write_bit(requeue)

        self.send_frame(MethodFrame(self.channel_id, 60, 100, args))

    def recover(self, requeue=False, cb=None):
        '''
        Ask server to redeliver all unacknowledged messages.
        '''
        args = Writer()
        args.write_bit(requeue)

        # The XML spec is incorrect; this method is always synchronous
        #  http://lists.rabbitmq.com/pipermail/rabbitmq-discuss/2011-January/010738.html
        self._recover_cb.append(cb)
        self.send_frame(MethodFrame(self.channel_id, 60, 110, args))
        self.channel.add_synchronous_cb(self._recv_recover_ok)

    def _recv_recover_ok(self, _method_frame):
        cb = self._recover_cb.popleft()
        if cb:
            cb()

    def _read_msg(self, method_frame, with_consumer_tag=False,
                  with_message_count=False):
        '''
        Support method to read a Message from the current frame buffer.
        Will return a Message, or re-queue current frames and raise a
        FrameUnderflow. Takes an optional argument on whether to read the
        consumer tag so it can be used for both deliver and get-ok.
        '''
        # No need to assert that is instance of Header or Content frames
        # because failure to access as such will result in exception that
        # channel will pick up and handle accordingly.
        header_frame = self.channel.next_frame()
        if header_frame:
            size = header_frame.size
            body = bytearray()
            rbuf_frames = deque([header_frame, method_frame])

            while len(body) < size:
                content_frame = self.channel.next_frame()
                if content_frame:
                    rbuf_frames.appendleft(content_frame)
                    body.extend(content_frame.payload.buffer())
                else:
                    self.channel.requeue_frames(rbuf_frames)
                    raise self.FrameUnderflow()
        else:
            self.channel.requeue_frames([method_frame])
            raise self.FrameUnderflow()

        if with_consumer_tag:
            consumer_tag = method_frame.args.read_shortstr()
        delivery_tag = method_frame.args.read_longlong()
        redelivered = method_frame.args.read_bit()
        exchange = method_frame.args.read_shortstr()
        routing_key = method_frame.args.read_shortstr()
        if with_message_count:
            message_count = method_frame.args.read_long()

        delivery_info = {
            'channel': self.channel,
            'delivery_tag': delivery_tag,
            'redelivered': redelivered,
            'exchange': exchange,
            'routing_key': routing_key,
        }
        if with_consumer_tag:
            delivery_info['consumer_tag'] = consumer_tag
        if with_message_count:
            delivery_info['message_count'] = message_count

        return Message(body=body, delivery_info=delivery_info,
                       **header_frame.properties)
