# -*- coding: utf-8 -*-

# Copyright (c) 2010-2012, GEM Foundation.
#
# OpenQuake is free software: you can redistribute it and/or modify it
# under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# OpenQuake 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 General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with OpenQuake.  If not, see <http://www.gnu.org/licenses/>.

"""
Classes dealing with amqp signalling between jobbers, workers and supervisors.
"""
import socket

import kombu
import kombu.entity
import kombu.messaging

from openquake.utils import config


def amqp_connect():
    """
    Connect to amqp broker with kombu using default configuration
    and return connection, channel and exchange as tuple.
    """
    cfg = config.get_section("amqp")
    connection = kombu.BrokerConnection(hostname=cfg['host'],
                                        userid=cfg['user'],
                                        password=cfg['password'],
                                        virtual_host=cfg['vhost'])
    channel = connection.channel()
    exchange = kombu.entity.Exchange(cfg['exchange'], type='topic',
                                     channel=channel)
    exchange.declare()
    return connection, channel, exchange


class AMQPMessageConsumer(object):
    """
    A baseclass to consume logging messages generated by an OpenQuake job.

    Typical use::

        class MyConsumer(AMQPMessageConsumer):
            def message_callback(self, payload, msg):
                # do something usefull
                pass

        MyConsumer('routing.key.#').run()
    """
    def __init__(self, routing_key, timeout=None):
        if timeout is not None:
            assert type(self).timeout_callback \
                   is not AMQPMessageConsumer.timeout_callback, \
                   "please override timeout_callback() method " \
                   "if you want to handle timeouts"
        self.timeout = timeout
        self.channel = self.connection = None
        self._stopped = False

        self.connection, self.channel, exchange = amqp_connect()
        queue = kombu.entity.Queue(exchange=exchange, channel=self.channel,
                                   routing_key=routing_key, exclusive=True)
        queue.queue_declare()
        queue.queue_bind()
        consumer = kombu.messaging.Consumer(
            self.channel, queue, auto_declare=False)
        consumer.register_callback(self._message_callback)
        consumer.consume()

    def run(self):
        """
        Consumer's main function of execution.
        """
        try:
            while not self._stopped:
                try:
                    self.connection.drain_events(timeout=self.timeout)
                except socket.timeout:
                    self._timeout_callback()
        finally:
            self.channel.close()
            self.connection.close()

    def stop(self):
        """
        Stop thread execution on the next loop iteration in :meth:`run`.
        """
        self._stopped = True

    def _message_callback(self, body, msg):  # pylint: disable=W0613
        """
        Run :meth:`message_callback` and handle :exc:`StopIteration` raised
        from there doing :meth:`stop`.
        """
        payload = msg.decode()
        try:
            self.message_callback(payload, msg)
        except StopIteration:
            msg.ack()
            self.stop()
        else:
            msg.ack()

    def _timeout_callback(self):
        """
        Run :meth:`timeout_callback` and handle :exc:`StopIteration` raised
        from there doing :meth:`stop`.
        """
        try:
            self.timeout_callback()
        except StopIteration:
            self.stop()

    def message_callback(self, payload, msg):
        """
        Called by :meth:`run` when a message is received.

        Can raise StopIteration to stop the loop inside :meth:`run`
        and terminate execution.
        """
        raise NotImplementedError()

    def timeout_callback(self):
        """
        Called by :meth:`run` each time the timeout expires.

        You need to implement this only if you specify a timeout when creating
        an instance of AMQPMessageConsumer.

        Can raise StopIteration to stop the loop inside :meth:`run`
        and terminate execution.
        """
        pass
