#!/usr/bin/env python

import sys
import os
import json
import pprint
import subprocess
import optparse
import shutil

ROOT_DIR = os.path.expanduser(os.path.join('~', '.hadoop-dev-env'))

DISTRIBUTION_MAP = {
    'cdh4.4.0': {
        'download_url_prefix': 'http://archive.cloudera.com/cdh4/cdh/4',
        'libs': {
            'hadoop': 'hadoop-2.0.0-cdh4.4.0',
            'hbase': 'hbase-0.94.6-cdh4.4.0',
            'zookeeper': 'zookeeper-3.4.5-cdh4.4.0',
        },
    },
    'cdh5.0.0': {
        'download_url_prefix': 'http://archive.cloudera.com/cdh5/cdh/5',
        'libs': {
            'cloudera search': 'cloudera-search-1.0.0-cdh5.0.0',
            'hadoop': 'hadoop-2.3.0-cdh5.0.0',
            'hbase': 'hbase-0.96.1.1-cdh5.0.0',
            'solr': 'solr-4.4.0-cdh5.0.0',
            'zookeeper': 'zookeeper-3.4.5-cdh5.0.0',
        },
    },
}

DEPENDENCY_MAP = {
    'hadoop': ['zookeeper'],
    'hbase': ['hadoop', 'zookeeper'],
    'cloudera-search': ['hadoop', 'zookeeper'],
    'zookeeper': [],
}

class HadoopDevException(Exception):
    pass

class UnknownCluster(HadoopDevException):
    def __str__(self):
        return 'unknown cluster: ' + self.message

class UnknownDistribution(HadoopDevException):
    def __str__(self):
        return 'unknown distribution: ' + self.message

class UnknownLibrary(HadoopDevException):
    def __str__(self):
        return 'unknown library: ' + self.message

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

class Cluster(object):
    def __init__(self, cluster_name):
        self.cluster_name = cluster_name

        # Setup a default config if one doesn't already exist.
        if not os.path.exists(ROOT_DIR):
            print 'creating a default ~/.hadoop-dev-env...'
            shutil.copytree(
                    os.path.join(
                        os.path.dirname(os.path.realpath(__file__)),
                        '..',
                        'default-config'),
                    ROOT_DIR)

        # load the configuration
        with open(os.path.join(ROOT_DIR, 'config')) as f:
            global_config = json.load(f)

        try:
            cluster_config = global_config[self.cluster_name]
        except KeyError:
            raise UnknownCluster, cluster_name

        # Extract our configuration
        self.local = cluster_config.get('local', False)
        self.version = cluster_config['version']
        self.config_dir = os.path.join(ROOT_DIR, 'conf', cluster_name)
       
        self.downloads_dir = os.path.join(ROOT_DIR, 'downloads')
        self.working_dir = os.path.join(ROOT_DIR, 'versions')

        # Look up info about the distribution.
        try:
            self.distribution = DISTRIBUTION_MAP[self.version]
        except KeyError:
            raise UnknownDistribution, self.version

        self.download_url_prefix = self.distribution['download_url_prefix']

        # Extract our libraries.
        self.libs = {}

        for lib_name in cluster_config['libs']:
            try:
                full_name = self.distribution['libs'][lib_name]
            except KeyError:
                raise UnknownLibrary, lib_name

            if lib_name == 'hadoop':
                self.libs[lib_name] = HadoopLibrary(self, full_name)
            elif lib_name == 'hbase':
                self.libs[lib_name] = HBaseLibrary(self, full_name)
            elif lib_name == 'cloudera_search':
                self.libs[lib_name] = ClouderaSearchLibrary(self, full_name)
            elif lib_name == 'zookeeper':
                self.libs[lib_name] = ZookeeperLibrary(self, full_name)
            else:
                raise UnknownLibrary, lib_name

    def execute(self, *args, **kwargs):
        env = dict(os.environ)

        for lib in self.libs.itervalues():
            lib.setup_environment(env)

        kwargs['env'] = env

        if kwargs.pop('verbose', False):
            print '% {}'.format(' '.join(args))

        subprocess.check_call(list(args), **kwargs)

    def _resolve_dependencies(self):
        # we need to start hadoop first.
        order = []

        def add(lib):
            if lib in order:
                return

            for dep_name in DEPENDENCY_MAP.get(lib.name, []):
                try:
                    child_lib = self.libs[dep_name]
                except KeyError:
                    raise HadoopDevException, 'missing dependency %s' % dep_name

                add(child_lib)

            order.append(lib)

        for lib in self.libs.itervalues():
            add(lib)

        return order

    def start(self):
        if not self.local:
            raise HadoopDevException, 'cannot start remote clusters'

        for lib in self._resolve_dependencies():
            lib.start()

    def stop(self):
        if not self.local:
            raise HadoopDevException, 'cannot stop remote clusters'

        for lib in reversed(self._resolve_dependencies()):
            lib.stop()

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

class Library(object):
    _name = None

    def __init__(self, cluster, full_name):
        self.cluster = cluster
        self.full_name = full_name

        self.working_dir = os.path.join(self.cluster.working_dir, self.full_name)

    def extract_working_dir(self):
        # Only extract tarball if we haven't already.
        if not os.path.exists(self.working_dir):
            download_path = self.download_tarball()

            # Make sure the working directory exists.
            if not os.path.exists(self.cluster.working_dir):
                os.makedirs(self.cluster.working_dir)

            # Extract to the working dir
            print 'extracting', download_path

            subprocess.check_call([
                'tar',
                '-xvf',
                download_path,
            ], cwd=os.path.dirname(self.working_dir))

        return self.working_dir

    def download_tarball(self):
        tarball = self.full_name + '.tar.gz'
        download_path = os.path.join(self.cluster.downloads_dir, tarball)

        # Only download tarball if we haven't already.
        if not os.path.exists(download_path):
            if not os.path.exists(self.cluster.downloads_dir):
                os.makedirs(self.cluster.downloads_dir)

            download_url = self.cluster.download_url_prefix + '/' + tarball

            print 'downloading', download_url

            subprocess.check_call([
                'curl',
                '-O',
                download_url,
            ], cwd=self.cluster.downloads_dir)

        return download_path

    def setup_environment(self, env):
        self.extract_working_dir()

        env['PATH'] = os.pathsep.join([
            os.path.join(self.working_dir, 'bin'),
            os.path.join(self.working_dir, 'sbin'),
        ] + env['PATH'].split(os.pathsep))

    def execute(self, *args, **kwargs):
        self.cluster.execute(*args, **kwargs)

    def start(self):
        raise NotImplementedError

    def stop(self):
        raise NotImplementedError

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

class HadoopLibrary(Library):
    name = 'hadoop'

    def setup_environment(self, env):
        super(HadoopLibrary, self).setup_environment(env)

        env['HADOOP_CONF_DIR'] = os.path.join(self.cluster.config_dir, 'hadoop-conf')
        env['YARN_CONF_DIR'] = os.path.join(self.cluster.config_dir, 'yarn-conf')

        env['HADOOP_PREFIX'] = self.working_dir
        env['HADOOP_HOME'] = self.working_dir
        env['HADOOP_COMMON_HOME'] = self.working_dir
        env['HADOOP_HDFS_HOME'] = self.working_dir
        env['HADOOP_MAPRED_HOME'] = self.working_dir
        env['HADOOP_YARN_HOME'] = self.working_dir

        if self.cluster.version == 'cdh5.0.0':
            env['CDH_MR2_HOME'] = self.working_dir

    def start(self):
        env = dict(os.environ)
        self.setup_environment(env)

        self.execute('start-dfs.sh', env=env)
        self.execute('start-yarn.sh', env=env)
        self.execute('yarn-daemons.sh', 'start', 'proxyserver', env=env)
        self.execute('mr-jobhistory-daemon.sh', 'start', 'historyserver', env=env)

    def stop(self):
        env = dict(os.environ)
        self.setup_environment(env)

        self.execute('mr-jobhistory-daemon.sh', 'stop', 'historyserver', env=env)
        self.execute('yarn-daemons.sh', 'stop', 'proxyserver', env=env)
        self.execute('stop-yarn.sh', env=env)
        self.execute('stop-dfs.sh', env=env)

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

class HBaseLibrary(Library):
    name = 'hbase'

    def setup_environment(self, env):
        super(HBaseLibrary, self).setup_environment(env)

        env['HBASE_CONF_DIR'] = os.path.join(self.cluster.config_dir, 'hbase-conf')
        env['HBASE_HOME'] = self.working_dir

    def start(self):
        env = dict(os.environ)
        self.setup_environment(env)

        self.execute('start-hbase.sh', env=env)

    def stop(self):
        env = dict(os.environ)
        self.setup_environment(env)

        self.execute('stop-hbase.sh', env=env)

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

class ClouderaSearchLibrary(Library):
    name = 'cloudera-search'

    def setup_environment(self, env):
        super(ClouderaSearchLibrary, self).setup_environment(env)

        env['SOLR_CONF_DIR'] = os.path.join(self.cluster.config_dir, 'search-conf')
        env['SOLR_HOME'] = self.working_dir
        env['SOLR_ZK_ENSEMBLE'] = cluster['zk_ensemble']
        env['SOLR_HDFS_HOME'] = cluster['hdfs_home']

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

class ZookeeperLibrary(Library):
    name = 'zookeeper'

    def setup_environment(self, env):
        super(ZookeeperLibrary, self).setup_environment(env)

        env['ZOOCFGDIR'] = os.path.join(self.cluster.config_dir, 'zookeeper-conf')
        env['ZOOKEEPER_HOME'] = self.working_dir

    def start(self):
        env = dict(os.environ)
        self.setup_environment(env)

        self.execute('zkServer.sh', 'start', env=env)

    def stop(self):
        env = dict(os.environ)
        self.setup_environment(env)

        self.execute('zkServer.sh', 'stop', env=env)

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

def run(cluster, *args):
    if len(args) == 0:
        print >> sys.stderr, 'command not specified'
        return 1

    cluster.execute(*args)

def start(cluster, *args):
    cluster.start()

def stop(cluster):
    cluster.stop()

def main(argv):
    if len(argv) < 3:
        print >> sys.stderr, 'please specify a config and command'
        return 1

    command = argv[1]

    if command not in ['run', 'start', 'stop']:
        print >> sys.stderr, 'unknown command', command
        return 1

    cluster_name = argv[2]
    argv = argv[3:]

    cluster = Cluster(cluster_name)

    if command == 'run':
        return run(cluster, *argv)
    elif command == 'start':
        return start(cluster, *argv)
    elif command == 'stop':
        return stop(cluster, *argv)
    else:
        print >> sys.stderr, 'unknown command', command
        return 1

if __name__ == '__main__':
    try:
        sys.exit(main(sys.argv))
    except HadoopDevException, e:
        print >> sys.stderr, e
        sys.exit(1)
    except subprocess.CalledProcessError, e:
        print >> sys.stderr, e
        sys.exit(1)
