#!/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 sys, os
from optparse import OptionParser

from thrift.Thrift import *

from thrift.transport import TSocket
from thrift.transport import TTransport
from thrift.protocol import TBinaryProtocol

from fb303 import *
from fb303.ttypes import *

def service_ctrl(
                 command,
                 port,
                 trans_factory = None,
                 prot_factory = None):
    """
    service_ctrl is a generic function to execute standard fb303 functions

    @param command: one of stop, start, reload, status, counters, name, alive
    @param port: service's port
    @param trans_factory: TTransportFactory to use for obtaining a TTransport. Default is
                          TBufferedTransportFactory
    @param prot_factory: TProtocolFactory to use for obtaining a TProtocol. Default is
                         TBinaryProtocolFactory
    """

    if command in ["status"]:
        try:
            status = fb303_wrapper('status', port, trans_factory, prot_factory)
            status_details = fb303_wrapper('get_status_details', port, trans_factory, prot_factory)

            msg = fb_status_string(status)
            if (len(status_details)):
                msg += " - %s" % status_details
            print msg

            if (status == fb_status.ALIVE):
                return 2
            else:
                return 3
        except:
            print "Failed to get status"
            return 3

    # scalar commands
    if command in ["version","alive","name"]:
        try:
            result = fb303_wrapper(command,  port, trans_factory, prot_factory)
            print result
            return 0
        except:
            print "failed to get ",command
            return 3

    # counters
    if command in ["counters"]:
        try:
            counters = fb303_wrapper('counters',  port, trans_factory, prot_factory)
            for counter in counters:
                print "%s: %d" % (counter, counters[counter])
            return 0
        except:
            print "failed to get counters"
            return 3


    # Only root should be able to run the following commands
    if os.getuid() == 0:
        # async commands
        if command in ["stop","reload"] :
            try:
                fb303_wrapper(command, port, trans_factory, prot_factory)
                return 0
            except:
                print "failed to tell the service to ", command
                return 3
    else:
        if command in ["stop","reload"]:
            print "root privileges are required to stop or reload the service."
            return 4

    print "The following commands are available:"
    for command in ["counters","name","version","alive","status"]:
        print "\t%s" % command
    print "The following commands are available for users with root privileges:"
    for command in ["stop","reload"]:
        print "\t%s" % command


    return 0;


def fb303_wrapper(command, port, trans_factory = None, prot_factory = None):
    sock = TSocket.TSocket('localhost', port)

    # use input transport factory if provided
    if (trans_factory is None):
        trans = TTransport.TBufferedTransport(sock)
    else:
        trans = trans_factory.getTransport(sock)

    # use input protocol factory if provided
    if (prot_factory is None):
        prot = TBinaryProtocol.TBinaryProtocol(trans)
    else:
        prot = prot_factory.getProtocol(trans)

    # initialize client and open transport
    fb303_client = FacebookService.Client(prot, prot)
    trans.open()

    if (command == 'reload'):
        fb303_client.reinitialize()

    elif (command == 'stop'):
        fb303_client.shutdown()

    elif (command == 'status'):
        return fb303_client.getStatus()

    elif (command == 'version'):
        return fb303_client.getVersion()

    elif (command == 'get_status_details'):
        return fb303_client.getStatusDetails()

    elif (command == 'counters'):
        return fb303_client.getCounters()

    elif (command == 'name'):
        return fb303_client.getName()

    elif (command == 'alive'):
        return fb303_client.aliveSince()

    trans.close()


def fb_status_string(status_enum):
    if (status_enum == fb_status.DEAD):
        return "DEAD"
    if (status_enum == fb_status.STARTING):
        return "STARTING"
    if (status_enum == fb_status.ALIVE):
        return "ALIVE"
    if (status_enum == fb_status.STOPPING):
        return "STOPPING"
    if (status_enum == fb_status.STOPPED):
        return "STOPPED"
    if (status_enum == fb_status.WARNING):
        return "WARNING"


def main():

    # parse command line options
    parser = OptionParser()
    commands=["stop","counters","status","reload","version","name","alive"]

    parser.add_option("-c", "--command", dest="command", help="execute this API",
                      choices=commands, default="status")
    parser.add_option("-p","--port",dest="port",help="the service's port",
                      default=9082)

    (options, args) = parser.parse_args()
    status = service_ctrl(options.command, options.port)
    sys.exit(status)


if __name__ == '__main__':
    main()