'''
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.connection import Connection
from haigha.classes.basic_class import BasicClass
from haigha.classes.exchange_class import ExchangeClass
from haigha.classes.protocol_class import ProtocolClass
from haigha.writer import Writer
from haigha.frames.method_frame import MethodFrame


class RabbitConnection(Connection):

    '''
    A connection specific to RabbitMQ that supports its extensions.
    '''

    def __init__(self, **kwargs):
        '''
        Initialize the connection
        '''
        class_map = kwargs.get('class_map', {}).copy()
        class_map.setdefault(40, RabbitExchangeClass)
        class_map.setdefault(60, RabbitBasicClass)
        class_map.setdefault(85, RabbitConfirmClass)
        kwargs['class_map'] = class_map

        super(RabbitConnection, self).__init__(**kwargs)


class RabbitExchangeClass(ExchangeClass):

    '''
    Exchange class Rabbit extensions
    '''

    def __init__(self, *args, **kwargs):
        super(RabbitExchangeClass, self).__init__(*args, **kwargs)
        self.dispatch_map[31] = self._recv_bind_ok
        self.dispatch_map[51] = self._recv_unbind_ok

        self._bind_cb = deque()
        self._unbind_cb = deque()

    # I hate the code copying here. Probably a better solution, like
    # functools.
    def declare(self, exchange, type, passive=False, durable=False,
                auto_delete=True, internal=False, nowait=True,
                arguments=None, ticket=None, cb=None):
        """
        Declare the exchange.

        exchange - The name of the exchange to declare
        type - One of
        """
        nowait = nowait and self.allow_nowait() and not cb

        args = Writer()
        args.write_short(ticket or self.default_ticket).\
            write_shortstr(exchange).\
            write_shortstr(type).\
            write_bits(passive, durable, auto_delete, internal, nowait).\
            write_table(arguments or {})
        self.send_frame(MethodFrame(self.channel_id, 40, 10, args))

        if not nowait:
            self._declare_cb.append(cb)
            self.channel.add_synchronous_cb(self._recv_declare_ok)

    def bind(self, exchange, source, routing_key='', nowait=True,
             arguments={}, ticket=None, cb=None):
        '''
        Bind an exchange to another.
        '''
        nowait = nowait and self.allow_nowait() and not cb

        args = Writer()
        args.write_short(ticket or self.default_ticket).\
            write_shortstr(exchange).\
            write_shortstr(source).\
            write_shortstr(routing_key).\
            write_bit(nowait).\
            write_table(arguments or {})
        self.send_frame(MethodFrame(self.channel_id, 40, 30, args))

        if not nowait:
            self._bind_cb.append(cb)
            self.channel.add_synchronous_cb(self._recv_bind_ok)

    def _recv_bind_ok(self, _method_frame):
        '''Confirm exchange bind.'''
        cb = self._bind_cb.popleft()
        if cb:
            cb()

    def unbind(self, exchange, source, routing_key='', nowait=True,
               arguments={}, ticket=None, cb=None):
        '''
        Unbind an exchange from another.
        '''
        nowait = nowait and self.allow_nowait() and not cb

        args = Writer()
        args.write_short(ticket or self.default_ticket).\
            write_shortstr(exchange).\
            write_shortstr(source).\
            write_shortstr(routing_key).\
            write_bit(nowait).\
            write_table(arguments or {})
        self.send_frame(MethodFrame(self.channel_id, 40, 40, args))

        if not nowait:
            self._unbind_cb.append(cb)
            self.channel.add_synchronous_cb(self._recv_unbind_ok)

    def _recv_unbind_ok(self, _method_frame):
        '''Confirm exchange unbind.'''
        cb = self._unbind_cb.popleft()
        if cb:
            cb()


class RabbitBasicClass(BasicClass):

    '''
    Support Rabbit extensions to Basic class.
    '''

    def __init__(self, *args, **kwargs):
        super(RabbitBasicClass, self).__init__(*args, **kwargs)
        self.dispatch_map[80] = self._recv_ack
        self.dispatch_map[120] = self._recv_nack

        self._ack_listener = None
        self._nack_listener = None

        self._msg_id = 0
        self._last_ack_id = 0

    def set_ack_listener(self, cb):
        '''
        Set a callback for ack listening, to be used when the channel is
        in publisher confirm mode. Will be called with a single integer
        argument which is the id of the message as returned from publish().

        cb(message_id)
        '''
        self._ack_listener = cb

    def set_nack_listener(self, cb):
        '''
        Set a callbnack for nack listening, to be used when the channel is
        in publisher confirm mode. Will be called with a single integer
        argument which is the id of the message as returned from publish()
        and a boolean flag indicating if it can be requeued.

        cb(message_id, reque)
        '''
        self._nack_listener = cb

    # Probably a better solution here, like functools
    def publish(self, *args, **kwargs):
        '''
        Publish a message. Will return the id of the message if publisher
        confirmations are enabled, else will return 0.
        '''
        if self.channel.confirm._enabled:
            self._msg_id += 1
        super(RabbitBasicClass, self).publish(*args, **kwargs)
        return self._msg_id

    def _recv_ack(self, method_frame):
        '''Receive an ack from the broker.'''
        if self._ack_listener:
            delivery_tag = method_frame.args.read_longlong()
            multiple = method_frame.args.read_bit()
            if multiple:
                while self._last_ack_id < delivery_tag:
                    self._last_ack_id += 1
                    self._ack_listener(self._last_ack_id)
            else:
                self._last_ack_id = delivery_tag
                self._ack_listener(self._last_ack_id)

    def nack(self, delivery_tag, multiple=False, requeue=False):
        '''Send a nack to the broker.'''
        args = Writer()
        args.write_longlong(delivery_tag).\
            write_bits(multiple, requeue)

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

    def _recv_nack(self, method_frame):
        '''Receive a nack from the broker.'''
        if self._nack_listener:
            delivery_tag = method_frame.args.read_longlong()
            multiple, requeue = method_frame.args.read_bits(2)
            if multiple:
                while self._last_ack_id < delivery_tag:
                    self._last_ack_id += 1
                    self._nack_listener(self._last_ack_id, requeue)
            else:
                self._last_ack_id = delivery_tag
                self._nack_listener(self._last_ack_id, requeue)


class RabbitConfirmClass(ProtocolClass):

    '''
    Implementation of Rabbit's confirm class.
    '''

    def __init__(self, *args, **kwargs):
        super(RabbitConfirmClass, self).__init__(*args, **kwargs)
        self.dispatch_map = {
            11: self._recv_select_ok,
        }

        self._enabled = False
        self._select_cb = deque()

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

    def select(self, nowait=True, cb=None):
        '''
        Set this channel to use publisher confirmations.
        '''
        nowait = nowait and self.allow_nowait() and not cb

        if not self._enabled:
            self._enabled = True
            self.channel.basic._msg_id = 0
            self.channel.basic._last_ack_id = 0
            args = Writer()
            args.write_bit(nowait)

            self.send_frame(MethodFrame(self.channel_id, 85, 10, args))

            if not nowait:
                self._select_cb.append(cb)
                self.channel.add_synchronous_cb(self._recv_select_ok)

    def _recv_select_ok(self, _method_frame):
        cb = self._select_cb.popleft()
        if cb:
            cb()
