#!/usr/bin/env python

#
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
#

import optparse, sys, time, os, re
from qpid.messaging import Connection
from qpid.messaging import Message as QpidMessage
from qpidtoollibs.broker import BrokerAgent
try:
    from uuid import uuid4
except ImportError:
    from qpid.datatypes import uuid4

# QMF address for the HA broker object.
HA_BROKER = "org.apache.qpid.ha:habroker:ha-broker"

class ExitStatus(Exception):
    """Raised if a command want's a non-0 exit status from the script"""
    def __init__(self, status): self.status = status

class Command:
    commands = {}

    def add(self, optname, metavar, type, help):
        self.op.add_option(optname, metavar=metavar, type=type, help=help, action="store")

    def __init__(self, name, help, arg_names=[]):
        Command.commands[name] = self
        self.name = name
        self.arg_names = arg_names
        usage="%s [options] %s\n\n%s"%(name, " ".join(arg_names), help)
        self.help = help
        self.op=optparse.OptionParser(usage)
        self.op.add_option("--sasl-mechanism", action="store", type="string", metavar="<mech>", help="SASL mechanism for authentication (e.g. EXTERNAL, ANONYMOUS, PLAIN, CRAM-MD, DIGEST-MD5, GSSAPI). SASL automatically picks the most secure available mechanism - use this option to override.")
        self.op.add_option("--ssl-certificate", action="store", type="string", metavar="<cert>", help="Client SSL certificate (PEM Format)")
        self.op.add_option("--ssl-key", action="store", type="string", metavar="<key>", help="Client SSL private key (PEM Format)")
        self.op.add_option("-b", "--broker", action="store", type="string", default="localhost:5672", metavar="<address>", help="Address of qpidd broker with syntax: [username/password@] hostname | ip-address [:<port>]")

    def connect(self, opts):
        conn_options = {}
        if opts.sasl_mechanism:
            conn_options['sasl_mechanisms'] = opts.sasl_mechanism
        if opts.ssl_certificate:
            conn_options['ssl_certfile'] = opts.ssl_certificate
        if opts.ssl_key:
            if not opts.ssl_certificate:
                self.op.error("missing '--ssl-certificate' (required by '--ssl-key')")
            conn_options['ssl_keyfile'] = opts.ssl_key
        conn_options['client_properties'] = {'qpid.ha-admin' : 1}
        connection = Connection.establish(opts.broker, **conn_options)
        qmf_broker = BrokerAgent(connection)
        ha_broker = qmf_broker.getHaBroker()
        return (connection, qmf_broker, ha_broker)

    def execute(self, args):
        opts, args = self.op.parse_args(args)
        if len(args) != len(self.arg_names)+1:
            self.op.print_help()
            raise Exception("Wrong number of arguments")
        connection, qmf_broker, ha_broker = self.connect(opts)
        if not ha_broker: raise Exception("HA module is not loaded on broker at %s" % opts.broker)
        try: self.do_execute(qmf_broker, ha_broker, opts, args)
        finally: connection.close()

    def do_execute(self, qmf_broker, opts, args):
        raise Exception("Command '%s' is not yet implemented"%self.name)

class PromoteCmd(Command):
    def __init__(self):
        Command.__init__(self, "promote","Promote broker from backup to primary")
    def do_execute(self, qmf_broker, ha_broker, opts, args):
        qmf_broker._method("promote", {}, HA_BROKER)
PromoteCmd()

class StatusCmd(Command):
    def __init__(self):
        Command.__init__(self, "status", "Print HA status")
        self.op.add_option(
            "--expect", type="string", metavar="<status>",
            help="Don't print status. Return 0 if it matches <status>, 1 otherwise")
        self.op.add_option(
            "--is-primary", action="store_true", default=False,
            help="Don't print status. Return 0 if the broker is primary, 1 otherwise")
        self.op.add_option(
            "--all", action="store_true", default=False,
            help="Print status for all brokers in the cluster.")
    def do_execute(self, qmf_broker, ha_broker, opts, args):
        if opts.is_primary:
            if not ha_broker.status in ["active", "recovering"]: raise ExitStatus(1)
        if opts.expect:
            if opts.expect != ha_broker.status: raise ExitStatus(1)
        if opts.all:
            opts.all=False
            brokers = re.sub(r'(^amqp:)|(tcp:)', "", ha_broker.brokersUrl).split(",")
            for b in brokers:
                opts.broker = b
                try:
                    connection, qmf_broker, ha_broker = self.connect(opts)
                    print b, ha_broker.status
                except Exception,e:
                    print b, e
        else:
            print ha_broker.status

StatusCmd()

class ReplicateCmd(Command):
    def __init__(self):
        Command.__init__(self, "replicate", "Set up replication from <queue> on <remote-broker> to <queue> on the current broker.", ["<queue>", "<remote-broker>"])
    def do_execute(self, qmf_broker, ha_broker, opts, args):
        qmf_broker._method("replicate", {"broker":args[1], "queue":args[2]}, HA_BROKER)
ReplicateCmd()

class SetCmd(Command):
    def __init__(self):
        Command.__init__(self, "set", "Set HA configuration settings")
        self.add("--brokers-url", "<url>", "string", "URL with address of each broker in the cluster. Used by brokers to connect to each other.")
        self.add("--public-url", "<url>", "string", "URL advertised to clients to connect to the cluster. May be a list or a VIP.")
        self.add("--backups", "<n>", "int", "Expect <n> backups to be running"),

    def do_execute(self, qmf_broker, ha_broker, opts, args):
        if (opts.brokers_url): qmf_broker._method("setBrokersUrl", {"url":opts.brokers_url}, HA_BROKER)
        if (opts.public_url): qmf_broker._method("setPublicUrl", {"url":opts.public_url}, HA_BROKER)
        if (opts.backups): qmf_broker._method("setExpectedBackups", {"expectedBackups":opts.backups}, HA_BROKER)

SetCmd()

class QueryCmd(Command):
    def __init__(self):
        Command.__init__(self, "query", "Print HA configuration and status")

    def do_execute(self, qmf_broker, ha_broker, opts, args):
        hb = ha_broker
        for x in [("Status:", hb.status),
                  ("Brokers URL:", hb.brokersUrl),
                  ("Public URL:", hb.publicUrl),
                  ("Expected Backups:", hb.expectedBackups),
                  ("Replicate: ", hb.replicateDefault)
                  ]:
            print "%-20s %s"%(x[0], x[1])

QueryCmd()

def print_usage(prog):
    print "usage: %s <command> [<arguments>]\n\nCommands are:\n"%prog
    for name, command in Command.commands.iteritems():
        help = command.help
        print "  %-12s %s."%(name, help.split(".")[0])
    print "\nFor help with a command type: %s <command> --help\n"%prog

def find_command(args):
    """Find a command among the arguments and options"""
    for arg in args:
        if arg in Command.commands:
            return Command.commands[arg]
    return None

def main_except(argv):
    """This version of main raises exceptions"""
    args=argv[1:]
    if args and args[0] == "--help-all":
        for c in Command.commands.itervalues():
            c.op.print_help(); print
    else:
        command = find_command(args)
        if not command:
            print_usage(os.path.basename(argv[0]));
            raise Exception("Command not found")
        command.execute(args)

def main(argv):
    try:
        main_except(argv)
        return 0
    except ExitStatus, e:
        return e.status
    except Exception, e:
        print e
        return 1

if __name__ == "__main__":
    sys.exit(main(sys.argv))