###############################################################################
# (c) Copyright 2017 CERN                                                     #
#                                                                             #
# This software is distributed under the terms of the GNU General Public      #
# Licence version 3 (GPL Version 3), copied verbatim in the file "COPYING".   #
#                                                                             #
# In applying this licence, CERN does not waive the privileges and immunities #
# granted to it by virtue of its status as an Intergovernmental Organization  #
# or submit itself to any jurisdiction.                                       #
###############################################################################
'''
Module grouping the common build functions.
'''

import datetime
import json
import logging
import os
import pika
import uuid

from collections import namedtuple
__author__ = 'Ben Couturier <ben.couturier@cern.ch>'


# Utility class to ease the use of the messages
Message = namedtuple('Message', ['body', 'method_frame', 'header_frame'],
                     rename=True)


# Our own exception to signal that the processing was interrupted and
# we shouldn't ack the message
class InterruptedProcessing(Exception):
    pass


# Various utilities for connection management
def _get_pwd_from_sys():
    """
    Get the RabbitMQ password from the environment of from a file on disk
    """
    # First checking the environment
    res = os.environ.get("RMQPWD", None)

    # Checking for the password in $HOME/private/rabbitmq.txt
    if res is None:
        fname = os.path.join(os.environ["HOME"], "private", "rabbitmq.txt")
        if os.path.exists(fname):
            with open(fname, "r") as f:
                data = f.readlines()
                if len(data) > 0:
                    res = data[0].strip()

    # Separate the username/password
    (username, password) = res.split("/")
    return username, password


def get_connection(host=None,
                   username=None, passwd=None,
                   port=5671, vhost='/lhcb', use_ssl=True):
    """
    Create a connection to RabbitMQ
    """

    envhost = os.environ.get("RMQHOST", None)
    if envhost:
        host = envhost
    if host is None:
        host = "lbmessagingbroker.cern.ch"

    envvhost = os.environ.get("RMQVHOST", None)
    if envvhost:
        vhost = envvhost

    if os.environ.get("RMQNOSSL", None):
        use_ssl = False

    if os.environ.get("RMQPORT", None):
        port = int(os.environ.get("RMQPORT"))

    if username is None or passwd is None:
        (username, passwd) = _get_pwd_from_sys()
    credentials = pika.PlainCredentials(username, passwd)
    cpar = " ".join([str(x) for x in [host, use_ssl, port, vhost]])
    logging.debug("Connection paremeters: %s" % cpar)
    params = pika.ConnectionParameters(host,
                                       ssl=use_ssl,
                                       port=port,
                                       virtual_host=vhost,
                                       credentials=credentials)
    return pika.BlockingConnection(params)


def check_channel(channel, connection=None):
    """
    Ugly hack to keep backwards compatibility
    We should receive a pika.Channel, but if we are given a
    pika.BlockingConnection, we just create a new channel from it....
    """
    checked_channel = None
    if connection is not None:
        checked_channel = connection.channel()
    else:
        if isinstance(channel, pika.BlockingConnection):
            checked_channel = channel.channel()
        else:
            checked_channel = channel
    return checked_channel


class QueueHelper(object):
    """
    Utility class to access the rabbitmq broker
    """

    def __init__(self, channel=None, connection=None):
        """
        Initialize the messenger class
        """
        self._channel = check_channel(channel, connection)
        import lbmessaging
        self.MAX_PRIORITY = lbmessaging.priority(lbmessaging.EMERGENCY)

    def _setup_exchange(self, channel, exchange_name, auto_delete=False):
        """
        Creates an exchange with the desired exchange_type
        """
        channel.exchange_declare(exchange=exchange_name,
                                 exchange_type='topic',
                                 auto_delete=auto_delete,
                                 durable=True)

    def basic_publish(self, exchange_name, routing_key, body, priority=1,
                      max_retry=3, headers=None, auto_delete=False,
                      extra_properties={}):
        """
        Send a message to the topic defined for the builds
        """
        channel = self._channel
        self._setup_exchange(channel, exchange_name,
                             auto_delete=auto_delete)
        now = datetime.datetime.utcnow().isoformat()
        ipd = now
        if headers and headers.get('initial_publish_date'):
            ipd = headers.get('initial_publish_date')
        props = pika.BasicProperties(delivery_mode=2, priority=priority,
                                     headers={'max_retry': max_retry,
                                              'initial_publish_date': ipd,
                                              # make message persistent
                                              'publish_date': now})
        for key, value in extra_properties.items():
            setattr(props, key, value)
        channel.basic_publish(exchange=exchange_name,
                              routing_key=routing_key,
                              body=body,
                              properties=props)

    def bind_queue_to_exchange(self, queue_name, exchange_name,
                               routing_keys=None,
                               route_only_to_queue=False,
                               auto_delete=False):
        """
        Binds the queue to an exchange, make sure the queue is persistent
        """

        return self._setup_queue(self._channel, exchange_name,
                                 queue_name, routing_keys,
                                 route_only_to_queue=route_only_to_queue,
                                 auto_delete=auto_delete)

    def _setup_queue(self, channel, exchange_name, queue_name=None,
                     routing_keys=None, route_only_to_queue=False,
                     auto_delete=False):
        """
        Setup a queue and binds
        """
        self._setup_exchange(channel, exchange_name,
                             auto_delete=auto_delete)
        if queue_name is None:
            # Anonymous queue is NOT persistent
            logging.info("Declaring anonymous queue")
            result = channel.queue_declare(exclusive=True,
                                           auto_delete=auto_delete,
                                           arguments={
                                               "x-max-priority":
                                                   self.MAX_PRIORITY})
            queue_name = result.method.queue
        else:
            # Named queues are persistent...
            logging.info("Declaring durable queue: %s" % queue_name)
            result = channel.queue_declare(durable=1, queue=queue_name,
                                           auto_delete=auto_delete,
                                           arguments={
                                               "x-max-priority":
                                                   self.MAX_PRIORITY})

        # defaults to the wildcard for topic channels
        if routing_keys is None:
            # If the routing keys list is None, we could listen for all the
            # messages in the exchange using # or we could listen only for the
            # messages for this specific queue. As an example, in the RPC mode,
            # both the client and the server are using the same exchange, but
            # we need to separate the messages for the server queue ( requests)
            # from those for the client anonymous queue (replays)
            if route_only_to_queue:
                routing_keys = [queue_name]
            else:
                routing_keys = ["#"]

        # Now binding the queue to the topic
        for routing_key in routing_keys:
            logging.info("Binding %s to %s for '%s'" % (queue_name,
                                                        exchange_name,
                                                        routing_key))
            channel.queue_bind(exchange=exchange_name,
                               queue=queue_name,
                               routing_key=routing_key)
        return queue_name


class Sender(QueueHelper):
    """
    Class used by cvmfs-lhcbdev.cern.ch to communicate with the outside world
    """

    def __init__(self, exchange, *args, **kwargs):
        """
        Initialize props
        """
        super(Sender, self).__init__(*args, **kwargs)
        self._exchange_name = exchange

    def send_message(self, routing_key, payload, priority=1, max_retry=3,
                     auto_delete=False):
        """
        Sends a message to the exchange
        """
        body = json.dumps(payload)
        self.basic_publish(self._exchange_name, routing_key, body,
                           priority=priority, max_retry=max_retry,
                           auto_delete=auto_delete
                           )
        logging.info("Publish done to %s" % (self._exchange_name))

    def resend_message(self, message, keep_retry_count=False):
        """
        Resend a message based on a received header frame payload
        """
        exchange_name = message.method_frame.exchange
        routing_key = message.method_frame.routing_key

        bodyjson = json.dumps(message.body)
        new_max_retry = message.header_frame.headers.get('max_retry', 0)
        if not keep_retry_count and 'max_retry' in message.header_frame.\
                headers.keys():
            new_max_retry = message.header_frame.headers['max_retry'] - 1

        self.basic_publish(exchange_name, routing_key, bodyjson,
                           priority=message.header_frame.priority,
                           max_retry=new_max_retry,
                           headers=message.header_frame.headers)
        logging.info("resend done to %s" % (exchange_name))


class Receiver(QueueHelper):
    """
    Initializes a queue to receive messages
    """

    def __init__(self, exchange, queue, routing_keys=None, auto_delete=False,
                 route_only_to_queue=False,
                 *args, **kwargs):
        """
        Initialize props
        """
        super(Receiver, self).__init__(*args, **kwargs)
        self._exchange_name = exchange
        self._queue_name = queue
        self.bind_queue_to_exchange(
            queue, exchange_name=exchange,
            routing_keys=routing_keys,
            route_only_to_queue=route_only_to_queue,
            auto_delete=auto_delete)

    def receive_messages(self):
        """
        Receive all messages in the queue and returns the list
        """
        message_list = []
        while True:
            message = self.receive_message()
            if message is None:
                break
            message_list.append(message)
        return message_list

    def receive_message(self, runBeforeAck=None, noAck=False):
        """
        Receive a message from the queue
        :param runBeforeAck: A method to be invoked before sending the ACK back
                             to the message broker
        :return: the message itself
        """
        message = None
        channel = self._channel
        method_frame, header_frame, body = channel.basic_get(
            queue=self._queue_name)
        if method_frame is not None:
            try:
                t = json.loads(body.decode('UTF-8'))
            except:
                t = body
            message = Message(t, method_frame, header_frame)
            if runBeforeAck:
                runBeforeAck(message)
            if not noAck:
                channel.basic_ack(method_frame.delivery_tag)
        return message

    def consume_message(self, processor):
        """
        Receive a message from the queue
        :param processor: A method to be invoked for each message, it receives
        a message and the channel in parameter and should return True to
        acknowledge the message/ False  to avoid doing so.
        """
        channel = self._channel

        def wrapper(channel, method_frame, header_frame, body):
            t = json.loads(body.decode('UTF-8'))
            message = Message(t, method_frame, header_frame)
            try:
                rc = processor(message)
                if rc:
                    channel.basic_ack(method_frame.delivery_tag)
            except StopIteration as e:
                # This is the normal exit method, we should consume the last
                # message so that it does not stay
                # the queue
                channel.basic_ack(method_frame.delivery_tag)
                channel.stop_consuming()

        channel.basic_consume(wrapper, queue=self._queue_name)
        try:
            channel.start_consuming()
        except InterruptedProcessing:
            channel.stop_consuming()
        except Exception as e:
            channel.stop_consuming()
            raise e


class AdminHelper(QueueHelper):
    """
       Class used for the administration helpers
    """

    def __init__(self, exchange, *args, **kwargs):
        """
        Initialize props
        """
        super(AdminHelper, self).__init__(*args, **kwargs)
        self._exchange_name = exchange

    def change_messages(self, queue_name, message_changer,
                        dest_exchange=None):
        # First we declare the listening and sending brokers.
        # If we what to change the messages in the same queue, the sender
        # exchange is the same as the receiving one
        receive_broker = Receiver(self._exchange_name, queue_name,
                                  channel=self._channel)
        if dest_exchange:
            sender_broker = Sender(dest_exchange, channel=self._channel)
        else:
            sender_broker = Sender(self._exchange_name, channel=self._channel)
        messages = []
        # We need to look at all the messages in the queue
        while True:
            # We receive each message one by one, but we do not send neither
            # ACK nor NACK ( noAck = true)
            message = receive_broker.receive_message(noAck=True)
            if message is None:
                break
            # We perform the message filtering and change
            (modified, message) = message_changer(message)

            # The messages are added in a list of message tuples that include
            # delivery tag for ACK or NACK, if it was modified (thus NACK is
            # needed) and the new message if modified
            if modified:

                routing_key = message.method_frame.routing_key
                max_retry = message.header_frame.headers.get('max_retry', 3)
                priority = message.header_frame.priority
                logging.info("[Message %s] %s (P:%s) was changed!" % (
                    message.method_frame.delivery_tag,
                    message.body, priority))

                messages.append({
                    'delivery_tag': message.method_frame.delivery_tag,
                    'has_been_modified': modified,
                    'updated_message': {
                        'routing_key': routing_key,
                        'payload': message.body,
                        'priority': priority,
                        'max_retry': max_retry
                    }})

            else:
                logging.info("[Message %s] %s (P:%s) was not changed!" % (
                    message.method_frame.delivery_tag,
                    message.body, message.header_frame.priority))
                messages.append({
                    'delivery_tag': message.method_frame.delivery_tag,
                    'has_been_modified': modified,
                    'updated_message': None
                })
        # If the message was modified, send first the new message into the
        # correct exchange and then ack the old message. If no updates were
        # done, just send a NACK for the existing message to leave it in the
        # queue.
        for message in messages:
            try:
                if message['has_been_modified']:
                    updated_message = message['updated_message']
                    sender_broker.send_message(updated_message['routing_key'],
                                               updated_message['payload'],
                                               priority=updated_message[
                                                   'priority'],
                                               max_retry=updated_message[
                                                   'max_retry'])
                    logging.info(
                        "[Message-%s] Sending ACK!" % message['delivery_tag'])
                    self._channel.basic_ack(message['delivery_tag'])
                else:
                    logging.info(
                        "[Message-%s] Sending NACK!" % message['delivery_tag'])
                    self._channel.basic_nack(message['delivery_tag'])
            except Exception as e:
                logging.error(
                    "Error processing message %s : %s. Sending NACK!" % (
                        message['delivery_tag'], str(e)))
                self._channel.basic_nack(message['delivery_tag'])

    def list_messages_in_queue(self, queue_name):
        """Display the messages in the queue"""
        receive_broker = Receiver(self._exchange_name, queue_name,
                                  channel=self._channel)
        display = []
        messages = []
        try:
            while True:
                message = receive_broker.receive_message(noAck=True)
                if message is None:
                    break
                messages.append(message.method_frame.delivery_tag)

                routing_key = message.method_frame.routing_key
                max_retry = message.header_frame.headers.get('max_retry', None)
                priority = message.header_frame.priority
                display.append([routing_key, max_retry, priority] +
                               [str(x) for x in message.body])
        except Exception as e:
            logging.error("Error gathering messages info: %s" % e)
        finally:
            for message in messages:
                self._channel.basic_nack(message)
        return display


class RPCServer(Receiver):
    """
     RPC Server that is waiting for messages from a RPC Client on a given
     exchange and it replaying to the client anonymous queue.
     """

    def __init__(self, exchange, queue, routing_keys=None, auto_delete=False,
                 *args, **kwargs):
        super(RPCServer, self).__init__(exchange, queue, routing_keys,
                                        route_only_to_queue=True,
                                        auto_delete=auto_delete,
                                        *args, **kwargs)
        self.exchange = exchange
        self.auto_delete = auto_delete
        self.counter_messages = 0

    def receive_and_replay_rpc_message(self, processor, limit=None):
        def wrapper(message):
            self.counter_messages += 1
            replay_payload = json.dumps(processor(message))
            routing_key = message.header_frame.reply_to
            extra_properties = {
                'correlation_id': message.header_frame.correlation_id
            }
            self.basic_publish(self.exchange, routing_key, replay_payload,
                               priority=self.MAX_PRIORITY,
                               auto_delete=self.auto_delete,
                               extra_properties=extra_properties)
            if limit is not None and limit <= self.counter_messages:
                raise StopIteration
            return replay_payload is not None
        self.consume_message(wrapper)


class RPCClient(QueueHelper):
    """
    RPC Client that is sending messages to a RPC Server on a given exchange
    and it is waiting for the server replay on a anonymous queue.
    """

    def __init__(self, exchange, routing_keys=None, auto_delete=False,
                 *args, **kwargs):
        super(RPCClient, self).__init__(*args, **kwargs)
        self.auto_delete = auto_delete
        self.exchange_name = exchange
        self.corr_id = str(uuid.uuid4())
        self.queue_name = self.bind_queue_to_exchange(
            None, exchange_name=exchange,
            route_only_to_queue=True,
            routing_keys=routing_keys, auto_delete=auto_delete)
        self.counter_messages = 0

    def send_and_receive_rpc_message(self, routing_key, payload, processor,
                                     limit=None):
        body = json.dumps(payload)
        channel = self._channel
        extra_properties = {
            'correlation_id': self.corr_id,
            'reply_to': self.queue_name
        }

        def wrapper(channel, method_frame, header_frame, body):
            self.counter_messages += 1
            t = json.loads(body.decode('UTF-8'))
            message = Message(t, method_frame, header_frame)
            if self.corr_id != header_frame.correlation_id:
                channel.basic_ack(method_frame.delivery_tag)
            else:
                rc = processor(message)
                if rc:
                    channel.basic_ack(method_frame.delivery_tag)
                else:
                    channel.basic_nack(method_frame.delivery_tag)
            if limit is not None and limit <= self.counter_messages:
                channel.stop_consuming()

        channel.basic_consume(wrapper, queue=self.queue_name)
        self.basic_publish(self.exchange_name, routing_key, body,
                           priority=self.MAX_PRIORITY,
                           auto_delete=self.auto_delete,
                           extra_properties=extra_properties
                           )
        try:
            channel.start_consuming()
        except InterruptedProcessing:
            channel.stop_consuming()
        except Exception as e:
            channel.stop_consuming()
            raise e
