#!/usr/bin/env python
# -*- coding: utf-8 -*-
# modjkapi - API for modjk management
# Copyright © 2008-2011, Red Hat, Inc.
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program 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 General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
"""
Command line application for gracefully putting hosts into and out of
rotation in a modjk balancer.
"""

__docformat__ = 'restructuredtext'

import sys

from os import system
from time import sleep

from modjkapi import JKManagerBalancerObjectFactory


def get_host_map(datasource):
    """
    Simply turns the datasource data into a map for usage.

    :Parameters:
       - `datasource`: the factory to use.
    """
    # Get all hosts
    hosts = []
    for balancer in datasource.objects():
        hosts += balancer.get_attribute_from_objects('host')
    # Move to set and kill off dupes
    hosts = set(hosts)

    # Make a nice host map so we can iterate by host.
    host_map = {}
    for host in hosts:
        host_map[host] = []
        for balancer in datasource.objects():
            for worker in balancer.filter_workers_by(host=host):
                host_map[host].append(worker)
    return host_map


def main(jkmanage_xml_url, wait_time, action, hosts):
    """
    The main entry point.

    :Parameters:
       - `jkmanage_xml_url`: the url that will provide the api.
       - `wait_time`: how long to wait on sleeps.
       - `action`: enable or disable.
       - `hosts`: hosts to enable/disable.
    """
    # make action lower to make matching simpler
    action = action.lower()
    # Get the datasource ... in this case we want the balancers
    datasource = JKManagerBalancerObjectFactory(jkmanage_xml_url)

    # Get the host_map to make things easier on us.
    host_map = get_host_map(datasource)

    # Start turning stuff off doing it host by host
    for key in hosts:

        # We iterate over the host_map multiple times so we can disable
        # and poll faster. If we disable all, then start polling we can
        # let the servers work on bleeding off while we are still asking
        # for disables. By the time we poll hopefully they will be all
        # disabled and ready to go.

        if action == 'disable':
            # For each worker on the machine disable it
            for worker in host_map[key]:
                print "-- Disable %s (%s for %s)" % (worker.name,
                                                     worker.host,
                                                     worker.balancer['name'])
                worker.disable(False)
            # Wait for all the workers to become non busy
            for worker in host_map[key]:
                while int(worker.busy) != 0:
                    print "waiting for %s to become idle." % worker.name,
                    print " Sleeping %s seconds." % wait_time
                    print "from server %s -- %s (%s)" % (worker.host,
                                                         worker.busy,
                                                         worker.elected)
                    print "last update from server was at %s." % (
                        datasource.last_time_updated)
                    print ""
                    sleep(wait_time)
                    datasource.refresh()
                    hosts = get_host_map(datasource)

        elif action == 'enable':
            # Re-enable the workers
            for worker in host_map[key]:
                print "++ Enable %s (%s for %s)" % (worker.name,
                                                    worker.host,
                                                    worker.balancer['name'])
                worker.enable(False)


if __name__ == '__main__':
    # Run main if we are executed via command line
    try:
        URL = sys.argv[1]
        WAIT = int(sys.argv[2])
        ACTION = sys.argv[3]
        HOSTS = sys.argv[4:]
        # or go ahead and go ...
        main(URL, WAIT, ACTION, HOSTS)
    except KeyboardInterrupt:
        print "\nHeard CTRL+c. Exiting ..."
    except Exception, ex:
        print ex
        print ('Usage: jk-rotate URL SECONDS_TO_WAIT enable/disable HOSTS '
               'SERVER1..')
        print ('Example: jk-rotate "http://myserver/jkmanage?mime=xml" 10 '
               'enable myserver.dom anotherserver.dom')
        raise SystemExit(1)
