#!/usr/bin/env python

#   Copyright 2014 Cloudera, Inc.
#
#   Licensed 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 getpass
import optparse
import os
import signal
import sys
import toml

from cm_api.api_client import ApiResource
from supervisor import childutils

# ------------------------------------------------------------------------------

class HueServer(object):
    def __init__(self, address, state):
        self.address = address
        self.state = state


    def __hash__(self):
        return hash((self.address, self.state))


    def __str__(self):
        return '%s (%s)' % (self.address, self.state)


    def __eq__(self, other):
        return isinstance(other, HueServer) and \
                self.address == other.address and \
                self.state == other.state


    def __cmp__(self, other):
        return cmp(
                (self.address, self.state),
                (other.address, other.state))

# ------------------------------------------------------------------------------

class MonitorHue(object):
    def __init__(
            self,
            cm_client,
            stdin=sys.stdin,
            stdout=sys.stdout,
            stderr=sys.stderr):
        self.cm_client = cm_client
        self.stdin = stdin
        self.stdout = stdout
        self.stderr = stderr

        self.hue_servers = None
        self.listeners = []


    def run_forever(self):
        self.tick()

        while True:
            headers, payload = childutils.listener.wait(self.stdin, self.stdout)

            # ignore non-tick events.
            if not headers['eventname'].startswith('TICK'):
                childutils.listener.ok(self.stdout)
                continue

            try:
                self.tick()
            finally:
                childutils.listener.ok(self.stdout)


    def tick(self):
        """Update the load balancer if any Hue servers have been added or removed"""

        hue_servers = self.get_hue_servers()

        # Don't do anything if we've already processed this set of hue servers.
        if self.hue_servers == hue_servers:
            return

        self.hue_servers = hue_servers

        print >> sys.stderr, 'updating server list:'
        for server in sorted(hue_servers):
            print >> sys.stderr, '  %s' % server

        for listener in self.listeners:
            listener.update(hue_servers)


    def get_hue_servers(self):
        """Fetch all the known hue servers"""

        hue_servers = set()

        for cluster in self.cm_client.get_all_clusters():
            for service in cluster.get_all_services():
                if service.type == 'HUE':
                    for role in service.get_roles_by_type("HUE_SERVER"):
                        host = self.cm_client.get_host(role.hostRef.hostId)
                        hostname = host.hostname

                        config = role.get_config(view='full')
                        port = config['hue_http_port']
                        port = (port.value is None and port.default) or port.value

                        address = '%s:%s' % (hostname, port)

                        hue_servers.add(HueServer(address, role.roleState))

        return hue_servers

# ------------------------------------------------------------------------------

class ConfigListener(object):
    def __init__(self, rpc, config, process_names):
        self.rpc = rpc
        self.config = config
        self.process_names = process_names

        self.config_file = self.config['config_file']

        with open(self.config['config_template']) as f:
            self.config_template = f.read()

        with open(self.config['server_template']) as f:
            self.server_template = f.read()


    def update(self, hue_servers):
        processes = self.rpc.supervisor.getAllProcessInfo()

        config = self.expand_config_template(hue_servers)

        with open(self.config_file, 'w') as f:
            print >> f, config

        for process in processes:
            # Ignore down processes.
            if not process['pid']:
                continue

            # Ignore processes we aren't monitoring.
            if process['name'] not in self.process_names:
                continue

            self.reload_process_config(process)


    def expand_config_template(self, hue_servers):
        config_servers = []

        for index, server in enumerate(hue_servers):
            config_servers.append(self.expand_server_template(index, server))

        return self.config_template % {
            'servers': '\n'.join(config_servers)
        }


    def expand_server_template(self, index, server):
        raise NotImplementedError


    def reload_process(self, process):
        raise NotImplementedError

# ------------------------------------------------------------------------------

class HAProxyListener(ConfigListener):
    def expand_server_template(self, index, server):
        """Expand the HAProxy server line"""

        return self.server_template % {
            'index': index,
            'address': server.address,
        }


    def reload_process_config(self, process):
        """Let haproxy-wrapper know the config file has changed."""

        os.kill(process['pid'], signal.SIGUSR2)

# ------------------------------------------------------------------------------

class NginxListener(ConfigListener):
    def __init__(self, *args, **kwargs):
        super(NginxListener, self).__init__(*args, **kwargs)

        self.alias_file = self.config['alias_file']
        self.generate_alias_file()


    def find_static_files(self):
        paths = (
            '/opt/cloudera/parcels/CDH/lib/hue/build/static/',
            '/usr/lib/hue/build/static/',
            os.path.join(os.path.dirname(__file__), '../../../build/static/'),
        )

        for path in paths:
            if os.path.exists(path):
                return path
        else:
            raise Exception('could not find the Hue static files')


    def generate_alias_file(self):
        static_files = self.find_static_files()
        with open(self.alias_file, 'w') as f:
            print >> f, 'alias %s;' % static_files


    def expand_server_template(self, index, server):
        """Expand the Nginx server line"""

        config_server = self.server_template % {
            'address': server.address,
        }

        # Let nginx know if this backend is offline.
        if server.state == 'STOPPED':
            config_server += ' down'

        return config_server + ';'


    def reload_process_config(self, process):
        """Let nginx know the config file has changed."""

        os.kill(process['pid'], signal.SIGHUP)

# ------------------------------------------------------------------------------

def main(argv):
    parser = optparse.OptionParser()

    parser.add_option('-c', '--config_file',
                      help='config file')
    parser.add_option('--haproxy',
                      default=[],
                      type='str',
                      action='append',
                      help='monitor this haproxy process')
    parser.add_option('--nginx',
                      default=[],
                      type='str',
                      action='append',
                      help='monitor this haproxy process')

    options, args = parser.parse_args(argv[1:])

    if options.config_file is None:
        options.config_file = 'etc/hue-lb.toml'

    try:
        with open(options.config_file) as f:
            config = toml.load(f)
    except IOError, e:
        print >> sys.stderr, 'error opening %s: %s' % (options.config_file, e)
        return 1

    try:
        hostname = config['cloudera-manager']['hostname']
    except KeyError:
        hostname = 'localhost'

    try:
        username = config['cloudera-manager']['username']
    except KeyError:
        username = 'admin'

    try:
        password = config['cloudera-manager']['password']
    except KeyError:
        password = 'admin'

    if len(options.haproxy) == 0 and len(options.nginx) == 0:
        print >> sys.stderr, 'no processes specified'
        return 1

    client = ApiResource(hostname, username=username, password=password)

    monitor = MonitorHue(client)

    rpc = childutils.getRPCInterface(os.environ)

    if len(options.haproxy) != 0:
        listener = HAProxyListener(rpc, config['haproxy'], options.haproxy)
        monitor.listeners.append(listener)

    if len(options.nginx) != 0:
        listener = NginxListener(rpc, config['nginx'], options.nginx)
        monitor.listeners.append(listener)

    monitor.run_forever()

    return 0

# ------------------------------------------------------------------------------

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