"""
Copyright (c) 2008 Red Hat, Inc.

This file is part of the Qpid async store library msgstore.so.

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
 USA

The GNU Lesser General Public License is available in the file COPYING.
"""

import re
from brokertest import BrokerTest
from qpid.messaging import Empty
from qmf.console import Session

    
def store_args(store_dir = None):
    """Return the broker args necessary to load the async store"""
    assert BrokerTest.store_lib 
    if store_dir == None:
        return []
    return ["--store-dir", store_dir]

class Qmf:
    """
    QMF functions not yet available in the new QMF API. Remove this and replace with new API when it becomes available.
    """
    def __init__(self, broker):
        self.__session = Session()
        self.__broker = self.__session.addBroker("amqp://localhost:%d"%broker.port())

    def add_exchange(self, exchange_name, exchange_type, alt_exchange_name=None, passive=False, durable=False,
                     arguments = None):
        """Add a new exchange"""
        amqp_session = self.__broker.getAmqpSession()
        if arguments == None:
            arguments = {}
        if alt_exchange_name:
            amqp_session.exchange_declare(exchange=exchange_name, type=exchange_type,
                                          alternate_exchange=alt_exchange_name, passive=passive, durable=durable,
                                          arguments=arguments)
        else:
            amqp_session.exchange_declare(exchange=exchange_name, type=exchange_type, passive=passive, durable=durable,
                                          arguments=arguments)
    
    def add_queue(self, queue_name, alt_exchange_name=None, passive=False, durable=False, arguments = None):
        """Add a new queue"""
        amqp_session = self.__broker.getAmqpSession()
        if arguments == None:
            arguments = {}
        if alt_exchange_name:
            amqp_session.queue_declare(queue_name, alternate_exchange=alt_exchange_name, passive=passive,
                                       durable=durable, arguments=arguments)
        else:
            amqp_session.queue_declare(queue_name, passive=passive, durable=durable, arguments=arguments)
    
    def delete_queue(self, queue_name):
        """Delete an existing queue"""
        amqp_session = self.__broker.getAmqpSession()
        amqp_session.queue_delete(queue_name)

    def _query(self, name, _class, package, alt_exchange_name=None):
        """Qmf query function which can optionally look for the presence of an alternate exchange name"""
        try:
            obj_list = self.__session.getObjects(_class=_class, _package=package)
            found = False
            for obj in obj_list:
                if obj.name == name:
                    found = True
                    if alt_exchange_name != None:
                        alt_exch_list = self.__session.getObjects(_objectId=obj.altExchange)
                        if len(alt_exch_list) == 0 or alt_exch_list[0].name != alt_exchange_name:
                            return False
                    break
            return found
        except Exception:
            return False
                

    def query_exchange(self, exchange_name, alt_exchange_name=None):
        """Test for the presence of an exchange, and optionally whether it has an alternate exchange set to a known
        value."""
        return self._query(exchange_name, "exchange", "org.apache.qpid.broker", alt_exchange_name)
    
    def query_queue(self, queue_name, alt_exchange_name=None):
        """Test for the presence of an exchange, and optionally whether it has an alternate exchange set to a known
        value."""
        return self._query(queue_name, "queue", "org.apache.qpid.broker", alt_exchange_name)
    
    def queue_message_count(self, queue_name):
        """Query the number of messages on a queue"""
        queue_list = self.__session.getObjects(_class="queue", _name=queue_name)
        if len(queue_list):
            return queue_list[0].msgDepth
    
    def queue_empty(self, queue_name):
        """Check if a queue is empty (has no messages waiting)"""
        return self.queue_message_count(queue_name) == 0
    
    def close(self):
        self.__session.delBroker(self.__broker)
        self.__session = None


class StoreTest(BrokerTest):
    """
    This subclass of BrokerTest adds some convenience test/check functions
    """
    
    def _chk_empty(self, queue, receiver):
        """Check if a queue is empty (has no more messages)"""
        try:
            msg = receiver.fetch(timeout=0)
            self.assert_(False, "Queue \"%s\" not empty: found message: %s" % (queue, msg))
        except Empty:
            pass

    @staticmethod
    def make_message(msg_count, msg_size):
        """Make message content. Format: 'abcdef....' followed by 'msg-NNNN', where NNNN is the message count"""
        msg = "msg-%04d" % msg_count
        msg_len = len(msg)
        buff = ""
        if msg_size != None and msg_size > msg_len:
            for index in range(0, msg_size - msg_len):
                if index == msg_size - msg_len - 1:
                    buff += "-"
                else:
                    buff += chr(ord('a') + (index % 26))
        return buff + msg
    
    # Functions for formatting address strings
    
    @staticmethod
    def _fmt_csv(string_list, list_braces = None):
        """Format a list using comma-separation. Braces are optionally added."""
        if len(string_list) == 0:
            return ""
        first = True
        str_ = ""
        if list_braces != None:
            str_ += list_braces[0]
        for string in string_list:
            if string != None:
                if first:
                    first = False
                else:
                    str_ += ", "
                str_ += string
        if list_braces != None:
            str_ += list_braces[1]
        return str_
    
    def _fmt_map(self, string_list):
        """Format a map {l1, l2, l3, ...} from a string list. Each item in the list must be a formatted map
        element('key:val')."""
        return self._fmt_csv(string_list, list_braces="{}") 
    
    def _fmt_list(self, string_list):
        """Format a list [l1, l2, l3, ...] from a string list."""
        return self._fmt_csv(string_list, list_braces="[]") 
    
    def addr_fmt(self, node_name, **kwargs):
        """Generic AMQP to new address formatter. Takes common (but not all) AMQP options and formats an address
        string."""
        # Get keyword args
        node_subject = kwargs.get("node_subject")
        create_policy = kwargs.get("create_policy")
        delete_policy = kwargs.get("delete_policy")
        assert_policy = kwargs.get("assert_policy")
        mode = kwargs.get("mode")
        link = kwargs.get("link", False)
        link_name = kwargs.get("link_name")
        node_type = kwargs.get("node_type")
        durable = kwargs.get("durable", False)
        link_reliability = kwargs.get("link_reliability")
        x_declare_list = kwargs.get("x_declare_list", [])
        x_bindings_list = kwargs.get("x_bindings_list", [])
        x_subscribe_list = kwargs.get("x_subscribe_list", [])
        
        node_flag = not link and (node_type != None or durable or len(x_declare_list) > 0 or len(x_bindings_list) > 0)
        link_flag = link and (link_name != None or durable or link_reliability != None or len(x_declare_list) > 0 or
                             len(x_bindings_list) > 0 or len(x_subscribe_list) > 0)
        assert not (node_flag and link_flag)
        
        opt_str_list = []
        if create_policy != None:
            opt_str_list.append("create: %s" % create_policy)
        if delete_policy != None:
            opt_str_list.append("delete: %s" % delete_policy)
        if assert_policy != None:
            opt_str_list.append("assert: %s" % assert_policy)
        if mode != None:
            opt_str_list.append("mode: %s" % mode)
        if node_flag or link_flag:
            node_str_list = []
            if link_name != None:
                node_str_list.append("name: \"%s\"" % link_name)
            if node_type != None:
                node_str_list.append("type: %s" % node_type)
            if durable:
                node_str_list.append("durable: True")
            if link_reliability != None:
                node_str_list.append("reliability: %s" % link_reliability)
            if len(x_declare_list) > 0:
                node_str_list.append("x-declare: %s" % self._fmt_map(x_declare_list))
            if len(x_bindings_list) > 0:
                node_str_list.append("x-bindings: %s" % self._fmt_list(x_bindings_list))
            if len(x_subscribe_list) > 0:
                node_str_list.append("x-subscribe: %s" % self._fmt_map(x_subscribe_list))
            if node_flag:
                opt_str_list.append("node: %s" % self._fmt_map(node_str_list))
            else:
                opt_str_list.append("link: %s" % self._fmt_map(node_str_list))
        addr_str = node_name
        if node_subject != None:
            addr_str += "/%s" % node_subject
        if len(opt_str_list) > 0:
            addr_str += "; %s" % self._fmt_map(opt_str_list)
        return addr_str
    
    def snd_addr(self, node_name, **kwargs):
        """ Create a send (node) address"""
        # Get keyword args
        topic = kwargs.get("topic")
        topic_flag = kwargs.get("topic_flag", False)
        auto_create = kwargs.get("auto_create", True)
        auto_delete = kwargs.get("auto_delete", False)
        durable = kwargs.get("durable", False)
        exclusive = kwargs.get("exclusive", False)
        ftd_count = kwargs.get("ftd_count")
        ftd_size = kwargs.get("ftd_size")
        policy = kwargs.get("policy", "flow-to-disk")
        exchage_type = kwargs.get("exchage_type")
        
        create_policy = None
        if auto_create:
            create_policy = "always"
        delete_policy = None
        if auto_delete:
            delete_policy = "always"
        node_type = None
        if topic != None or topic_flag:
            node_type = "topic"
        x_declare_list = ["\"exclusive\": %s" % exclusive]
        if ftd_count != None or ftd_size != None:
            queue_policy = ["\'qpid.policy_type\': %s" % policy]
            if ftd_count:
                queue_policy.append("\'qpid.max_count\': %d" % ftd_count)
            if ftd_size:
                queue_policy.append("\'qpid.max_size\': %d" % ftd_size)
            x_declare_list.append("arguments: %s" % self._fmt_map(queue_policy))
        if exchage_type != None:
            x_declare_list.append("type: %s" % exchage_type)
            
        return self.addr_fmt(node_name, topic=topic, create_policy=create_policy, delete_policy=delete_policy,
                             node_type=node_type, durable=durable, x_declare_list=x_declare_list)
    
    def rcv_addr(self, node_name, **kwargs):
        """ Create a receive (link) address"""
        # Get keyword args
        auto_create = kwargs.get("auto_create", True)
        auto_delete = kwargs.get("auto_delete", False)
        link_name = kwargs.get("link_name")
        durable = kwargs.get("durable", False)
        browse = kwargs.get("browse", False)
        exclusive = kwargs.get("exclusive", False)
        binding_list = kwargs.get("binding_list", [])
        ftd_count = kwargs.get("ftd_count")
        ftd_size = kwargs.get("ftd_size")
        policy = kwargs.get("policy", "flow-to-disk")
        
        create_policy = None
        if auto_create:
            create_policy = "always"
        delete_policy = None
        if auto_delete:
            delete_policy = "always"
        mode = None
        if browse:
            mode = "browse" 
        x_declare_list = ["\"exclusive\": %s" % exclusive]
        if ftd_count != None or ftd_size != None:
            queue_policy = ["\'qpid.policy_type\': %s" % policy]
            if ftd_count:
                queue_policy.append("\'qpid.max_count\': %d" % ftd_count)
            if ftd_size:
                queue_policy.append("\'qpid.max_size\': %d" % ftd_size)
            x_declare_list.append("arguments: %s" % self._fmt_map(queue_policy))
        x_bindings_list = []
        for binding in binding_list:
            x_bindings_list.append("{exchange: %s, key: %s}" % binding)
        if durable: reliability = 'at-least-once'
        else: reliability = None
        return self.addr_fmt(node_name, create_policy=create_policy, delete_policy=delete_policy, mode=mode, link=True,
                             link_name=link_name, durable=durable, x_declare_list=x_declare_list,
                             x_bindings_list=x_bindings_list, link_reliability=reliability)
    
    def check_message(self, broker, queue, exp_msg, transactional=False, empty=False, ack=True, browse=False):
        """Check that a message is on a queue by dequeuing it and comparing it to the expected message"""
        return self.check_messages(broker, queue, [exp_msg], transactional, empty, ack, browse)
        
    def check_messages(self, broker, queue, exp_msg_list, transactional=False, empty=False, ack=True, browse=False,
                       emtpy_flag=False):
        """Check that messages is on a queue by dequeuing them and comparing them to the expected messages"""
        if emtpy_flag:
            num_msgs = 0
        else:
            num_msgs = len(exp_msg_list)
        ssn = broker.connect().session(transactional=transactional)
        rcvr = ssn.receiver(self.rcv_addr(queue, browse=browse), capacity=num_msgs)
        if num_msgs > 0:
            try:
                recieved_msg_list = [rcvr.fetch(timeout=0) for i in range(num_msgs)]
            except Empty:
                self.assert_(False, "Queue \"%s\" is empty, unable to retrieve expected message %d." % (queue, i))
            for i in range(0, len(recieved_msg_list)):
                self.assertEqual(recieved_msg_list[i].content, exp_msg_list[i].content)
                self.assertEqual(recieved_msg_list[i].correlation_id, exp_msg_list[i].correlation_id)
        if empty:
            self._chk_empty(queue, rcvr)
        if ack:
            ssn.acknowledge()
            if transactional:
                ssn.commit()
            ssn.connection.close()
        else:
            if transactional:
                ssn.commit()
            return ssn
            
    
    # Functions for finding strings in the broker log file (or other files)

    @staticmethod
    def _read_file(file_name):
        """Returns the content of file named file_name as a string"""
        file_handle = file(file_name)
        try:
            return file_handle.read()
        finally:
            file_handle.close()
    
    def _get_hits(self, broker, search):
        """Find all occurrences of the search in the broker log (eliminating possible duplicates from msgs on multiple
        queues)"""
        # TODO: Use sets when RHEL-4 is no longer supported
        hits = []
        for hit in search.findall(self._read_file(broker.log)):
            if hit not in hits:
                hits.append(hit) 
        return hits
    
    def _reconsile_hits(self, broker, ftd_msgs, release_hits):
        """Remove entries from list release_hits if they match the message id in ftd_msgs. Check for remaining
        release_hits."""
        for msg in ftd_msgs:
            found = False
            for hit in release_hits:
                if str(msg.id) in hit:
                    release_hits.remove(hit)
                    #print "Found %s in %s" % (msg.id, broker.log)
                    found = True
                    break
            if not found:
                self.assert_(False, "Unable to locate released message %s in log %s" % (msg.id, broker.log))
        if len(release_hits) > 0:
            err = "Messages were unexpectedly released in log %s:\n" % broker.log
            for hit in release_hits:
                err += "  %s\n" % hit
            self.assert_(False, err)
        
    def check_msg_release(self, broker, ftd_msgs):
        """ Check for 'Content released' messages in broker log for messages in ftd_msgs"""
        hits = self._get_hits(broker, re.compile("debug Message id=\"[0-9a-f-]{36}\"; pid=0x[0-9a-f]+: "
                                                 "Content released$", re.MULTILINE))
        self._reconsile_hits(broker, ftd_msgs, hits)
        
    def check_msg_release_on_commit(self, broker, ftd_msgs):
        """ Check for 'Content released on commit' messages in broker log for messages in ftd_msgs"""
        hits = self._get_hits(broker, re.compile("debug Message id=\"[0-9a-f-]{36}\"; pid=0x[0-9a-f]+: "
                                                 "Content released on commit$", re.MULTILINE))
        self._reconsile_hits(broker, ftd_msgs, hits)
        
    def check_msg_release_on_recover(self, broker, ftd_msgs):
        """ Check for 'Content released after recovery' messages in broker log for messages in ftd_msgs"""
        hits = self._get_hits(broker, re.compile("debug Message id=\"[0-9a-f-]{36}\"; pid=0x[0-9a-f]+: "
                                                 "Content released after recovery$", re.MULTILINE))
        self._reconsile_hits(broker, ftd_msgs, hits)
    
    def check_msg_block(self, broker, ftd_msgs):
        """Check for 'Content release blocked' messages in broker log for messages in ftd_msgs"""
        hits = self._get_hits(broker, re.compile("debug Message id=\"[0-9a-f-]{36}\"; pid=0x[0-9a-f]+: "
                                                 "Content release blocked$", re.MULTILINE))
        self._reconsile_hits(broker, ftd_msgs, hits)
     
    def check_msg_block_on_commit(self, broker, ftd_msgs):
        """Check for 'Content release blocked' messages in broker log for messages in ftd_msgs"""
        hits = self._get_hits(broker, re.compile("debug Message id=\"[0-9a-f-]{36}\"; pid=0x[0-9a-f]+: "
                                                 "Content release blocked on commit$", re.MULTILINE))
        self._reconsile_hits(broker, ftd_msgs, hits)
       
        
