# FIXME: debug flag, also prints out values every time subprocess
# is used!
"""osp-pipeline Cluster Manager

Start and deploy a cluster, run jobs on that cluster,
update the cluster, get logs.

Only works in Python 3.5 and later!

REQUIRES:
  * docopt! `pip3 install docopt==0.6.2`!
  * AWS CLI `pip install awscli`
  * `SLACK_API_TOKEN` environment variable
    See: https://api.slack.com/custom-integrations/legacy-tokens
  * `pip3 install slackclient==1.1.0`
  * `OSP_CLUSTER_NS` environment variable, e.g.,
    `export OSP_CLUSTER_NS='exampleName'` Alphanumeric only.

It's important that this script is run from the root
of the project directory, e.g.:

  `python3 deploy/cluster.py`

Usage:
    deploy/cluster.py deploy [--very-verbose]
    deploy/cluster.py update [--very-verbose]
    deploy/cluster.py masterui
    deploy/cluster.py terminate [--non-interactive]
    deploy/cluster.py logs <destination> [--very-verbose]
    deploy/cluster.py run <job> <logs_destination> [--very-verbose]
    deploy/cluster.py --version
    deploy/cluster.py (--help|-h)

Options
    -h --help         Show this screen.
    --version         Show version.
    --very-verbose    Enables ALL the verbose options.

"""

import os
import sys
import json
import time
import logging
import subprocess

from docopt import docopt
from slackclient import SlackClient

__version__ = '0.1.0'

NAMESPACE = os.environ['OSP_CLUSTER_NS']
"""str: this must be a string matching the one used in the ansible
playbooks for tagging to group clusters.
"""

SLEEPING_TO_LET_INSTANCES_INITIALIZE = '''
Sleep for {seconds} seconds to let the EC2 cluster instances
initialize before we deploy anything on them.

It sometimes takes up to {seconds} before the cluster is ready.
A better alternative to sleeping would be polling (a possible TODO).
'''

CLUSTER_EXPENSE_WARNING = '''
{namespace} cluster is now running! It's expensive to run, so
please terminate the cluster (its EC2 instances) as soon as
you are finished!

See README.md for more details!
'''

RUN_JOB_LOG_WARNING = '''
About to send {job} to {master_ip} via SSH. Make sure
you have SSH timeouts disabled on your system (see:
README.md).

You can check on the status of this job by looking at the
stdout for master stored in {logs_destination}/0/.
'''

SLACK_CLUSTER_STARTED_NOTICE = '''
🏁🆕 Starting {namespace} cluster deployent! ⏳⏲️  Please await further
details (this will take me a while)... 
'''

SLACK_CLUSTER_DEPLOYED_NOTICE =  '''
💚🆙 {namespace} cluster is up: {webui_address}
'''

SLACK_CLUSTER_TERMINATED_NOTICE = '''
💥🗑️ Terminated these EC2 instances belonging to the
{namespace} cluster: {terminated_instances}
'''

SLACK_CLUSTER_JOB_STARTED_NOTICE = '''
🏃▶️ Running {job} on the {namespace} cluster!
'''

SLACK_CLUSTER_JOB_ENDED_NOTICE = '''
🛑⏹️ Ended {job} on the {namespace} cluster!
'''

PATHS_TO_COPY_FROM_NODES = [
    '/tmp/spark-logs',
    '/var/log/syslog',
]
"""list: A list of paths which will be copied
from both the master and worker nodes (all nodes).

"""

NEVER_COPY_FILES_MATCHING = '*.jar'
"""str: files matching this pattern will not be copied
from spark nodes.

"""

BORDER_LEFT = '\u2588\u2593\u2592\u2591 '
BORDER_RIGHT = BORDER_LEFT[::-1]
LOG_FORMAT = (
    BORDER_LEFT
    + ' %(levelname)s @ %(asctime)s '
    + BORDER_RIGHT
    + '\n%(message)s'
)


def print_heading(one_line_of_text: str) -> str:
    print(BORDER_LEFT + one_line_of_text.upper() + BORDER_RIGHT)


def ansible_playbook(playbook: str, verbose=False) -> tuple:
    ansible_playbook_command = [
        'ansible-playbook',
        playbook,
    ]
    if verbose:
        ansible_playbook_command.append('-vvvv')
    process = subprocess.Popen(
        ansible_playbook_command,
    )
    output, error = process.communicate()
    return output, error


def ec2_refresh_cache(verbose=False) -> tuple:
    """Refresh Ansible's EC2 cache (which is handy for
    getting info on EC2 instances in the cluster!).

    Warning:
        Assumes you're in project root OR in `deploy/`!

    Arguments:
        verbose (optional): If True the output of the process
            will be printed.

    Returns:
        output, error: ...

    """

    if os.getcwd().endswith('deploy'):
        ec2_command = './ec2/ec2.py'
    else:
        ec2_command = './deploy/ec2/ec2.py'

    process = subprocess.Popen(
        [ec2_command, '--refresh-cache'],
        stdout=subprocess.PIPE,  # FIXME?
    )
    output, error = process.communicate()
    if verbose and any([output, error]):
        logging.debug(output)
        if error:
            logging.error(error)
    return output, error


def update_cluster(verbose=False) -> tuple:
    """Update a running Spark cluster with new code!

    Warning:
        Assumes you're in project root!

    Arguments:
        verbose: If True log output from command(s)!

    Returns:
        output, error: ...

    """

    os.chdir('deploy')
    output, error = ansible_playbook('cluster.deploy.yml', verbose=verbose)
    os.chdir('..')
    return output, error


def get_master_ip() -> str:
    """Return the IP of the Spark master node.

    Warning:
        See: ec2_refresh_cache()

    Raises:
        AssertionError: If there's more than one EC2 instance
            tagged with master, because then it cannot be told
            which is truly the master node.

    """

    ec2py_json = json.loads(ec2_refresh_cache()[0])
    spark_master_list = ec2py_json['tag_mit_%s_spark_master' % NAMESPACE]
    assert len(spark_master_list) == 1
    return spark_master_list[0]


def get_slack_client() -> SlackClient:
    """
    Get SLACK_API_TOKEN from
    https://api.slack.com/custom-integrations/legacy-tokens
    """
    slack_token = os.environ['SLACK_API_TOKEN']
    return SlackClient(slack_token)


def deploy_cluster(verbose=False) -> str:
    """Deploy a Spark cluster on AWS EC2.

    Assumes you're in project root!

    Returns:
        Spark master web interface address

    """

    # Post a nice Slack notification about deploy starting
    slack_client = get_slack_client()
    slack_client.api_call(
        'chat.postMessage',
        channel='#cluster-status',
        text=SLACK_CLUSTER_STARTED_NOTICE.format(namespace=NAMESPACE),
    )

    os.chdir('deploy')
    ec2_refresh_cache(verbose=verbose)

    # Run the cluster start Ansible playbook
    print_heading('start the ec2 instances')
    output, error = ansible_playbook('cluster.start.yml', verbose=verbose)

    # Let instances start up before we progress and refresh
    # the cache so we have correct node information
    sleep_for_x_seconds = 60
    logging.info(
        SLEEPING_TO_LET_INSTANCES_INITIALIZE.format(seconds=sleep_for_x_seconds),
    )
    time.sleep(sleep_for_x_seconds)
    # ... now hopefully all the instances have started!
    ec2_refresh_cache(verbose=verbose)

    # Run the cluster *deploy* Ansible playbook
    print_heading("configure the ec2 instances")
    output, error = ansible_playbook('cluster.deploy.yml', verbose=verbose)

    # Remind the user to shutdown their cluster!
    logging.warning(CLUSTER_EXPENSE_WARNING.format(namespace=NAMESPACE))

    # Return the Spark master web interface address
    os.chdir('..')
    spark_master_webui_address = 'http://%s:8080/' % get_master_ip()
    print_heading('Spark Master Web Interface Address')
    print(spark_master_webui_address)

    # Post a nice Slack notification and return...
    slack_client = get_slack_client()
    slack_client.api_call(
        'chat.postMessage',
        channel='#cluster-status',
        text=SLACK_CLUSTER_DEPLOYED_NOTICE.format(
            namespace=NAMESPACE,
            webui_address=spark_master_webui_address
        ),
    )
    return spark_master_webui_address


def fetch_logs(logs_destination: str, verbose=False) -> None:
    """Grab logs from all Spark nodes!

    Assumes you're in project root.

    If /tmp/some-dir-for-logs doesn't exist, this script will
    create that directory.

    """

    os.chdir('deploy')
    destination = logs_destination

    # Get the output of ./ec2/ec2.py --refresh-cache so we can find the
    # spark nodes by specific tags, then parse the JSON into a Python
    # dictionary!
    ec2py_json = json.loads(ec2_refresh_cache()[0])
    spark_workers = ec2py_json['tag_mit_%s_spark_worker' % NAMESPACE]
    spark_master = ec2py_json['tag_mit_%s_spark_master' % NAMESPACE]  # a list!

    for i, spark_node_ip in enumerate(spark_master + spark_workers):
        # Make sure the destination is ready!
        where_to_copy = os.path.join(destination, str(i))
        try:
            os.makedirs(where_to_copy)
        except os.error:
            pass

        # Loop through all the spark nodes, using rsync to copy each
        # file/directory in PATHS_TO_COPY_FROM_NODES
        host = 'ubuntu@%s' % spark_node_ip
        what_to_copy = list(
            map(lambda x: host + ':' + x, PATHS_TO_COPY_FROM_NODES),
        )
        for rsync_source in what_to_copy:
            subprocess.call([
                'rsync',
                '-rv',
                '-e',
                'ssh',
                '--exclude',
                NEVER_COPY_FILES_MATCHING,
                rsync_source,
                where_to_copy,
            ])

    print('Logs gathered from all Spark nodes and stored in %s!' % destination)
    print("The master node's logs are stored in %s!" % os.path.join(destination, '0'))
    os.chdir('..')


def terminate_spark_nodes(interactive: bool = True) -> None:
    """Delete the cluster's EC2 instances (the Spark nodes).

    Will ask and provide information before deleting.

    Arguments:
        interactive: When this is True keyboard confirmation
            is required to terminate nodes. Otherwise
            if False the nodes will be terminated without
            keyboard confirmation.

    """

    # First find all the worker and master nodes and display their
    # ec2 ID as well as their IP and when they were started.
    # ec2_launch_time

    ec2py_json = json.loads(ec2_refresh_cache()[0])
    spark_master_list = ec2py_json['tag_mit_%s_spark_master' % NAMESPACE]
    assert len(spark_master_list) == 1
    master_ip = spark_master_list[0]

    spark_worker_ips = ec2py_json['tag_mit_%s_spark_worker' % NAMESPACE]

    marked_for_deletion_info = []
    ids_marked_for_deletion = []
    for ip_address in spark_worker_ips + [master_ip]:
        node_info = ec2py_json['_meta']['hostvars'][ip_address]

        if ip_address == master_ip:
            node_type = 'master'
        else:
            node_type = 'worker'

        marked_for_deletion_info.append(
            {
                'type': node_type,
                'ip_address': ip_address,
                'ec2_id': node_info['ec2_id'],
                'launch_time': node_info['ec2_launch_time'],
            },
        )
        ids_marked_for_deletion.append(node_info['ec2_id'])

    for instance in marked_for_deletion_info:
        print(
            '''EC2 Instance ID: {ec2_id}
            Spark Node Type: {type}
            IP Address: {ip_address}
            Launch Time: {launch_time}
            '''.format(**instance)
        )

    print(
        'In summary, the EC2 instances with these IDs will be deleted: %s'
        % ', '.join(ids_marked_for_deletion)
    )
    print()

    if interactive:
        while True:
            answer = input('Terminate above EC2 instances (type "yes" or "no")? ').lower()
            if answer == 'yes':
                break
            elif answer == 'no':
                print('termination aborted')
                return False
            else:
                print('You must type a "yes" or a "no"')

    # now delete what we found since confirmed
    command = [
        'aws',
        'ec2',
        'terminate-instances',
        '--instance-ids',
    ]
    command = command + ids_marked_for_deletion
    subprocess.call(command)
    logging.info(
        'Successfully terminated: %s' % ', '.join(ids_marked_for_deletion),
    )
    slack_client = get_slack_client()
    slack_client.api_call(
        'chat.postMessage',
        channel='#cluster-status',
        text=SLACK_CLUSTER_TERMINATED_NOTICE.format(
            namespace=NAMESPACE,
            terminated_instances=', '.join(ids_marked_for_deletion),
        ),
    )


def run_job(job_to_run: str, logs_destination: str, verbose=False) -> None:
    """Run a spark job and bring the log from the spark-submit and
    the workers all to specified output.

    Must be ran from root...

    Arguments:
        job_to_run: e.g., catalog.jobs.load_otl
        logs_destination: where to store all the Spark node logs,
            e.g., /tmp/wowzers
        verbose (optional): Extra output!

    """

    # first create the directory which will contain the stdout
    # dumped live from running the SSH command...
    master_logs_directory = os.path.join(logs_destination, '0')
    try:
        os.makedirs(master_logs_directory)
    except OSError as e:
        # The directory already exists!
        logging.warning("%s already exists!" % master_logs_directory)

    # warning before running the job!
    master_ip = get_master_ip()
    logging.warning(
        RUN_JOB_LOG_WARNING.format(
            job=job_to_run,
            master_ip=master_ip,
            logs_destination=logs_destination,
        ),
    )

    # Slack notification about job starting
    slack_client = get_slack_client()
    slack_client.api_call(
        'chat.postMessage',
        channel='#cluster-status',
        text=SLACK_CLUSTER_JOB_STARTED_NOTICE.format(
            namespace=NAMESPACE,
            job=job_to_run,
        ),
    )

    # The incredibly complicated command! Wew!
    # FIXME: this string is just too long! break it into multiple lines!
    command_to_run_on_master =  'sudo docker exec -t $(sudo docker ps -aqf "name=osp-pipeline") spark-submit bin/run_job.py {job}'.format(
        job=job_to_run,
    )
    # FIXME: string too long!!!
    super_long_ssh_command = "ssh ubuntu@{master_ip} '{command_to_run_on_master}' >> {spark_submit_out}".format(
        master_ip=master_ip,
        command_to_run_on_master=command_to_run_on_master,
        spark_submit_out=os.path.join(master_logs_directory, 'spark-submit-out'),
    )
    process = subprocess.Popen(
        super_long_ssh_command,
        stdout=subprocess.PIPE,
        shell=True,
    )
    output, error = process.communicate()

    # Slack notification about job ending
    slack_client = get_slack_client()
    slack_client.api_call(
        'chat.postMessage',
        channel='#cluster-status',
        text=SLACK_CLUSTER_JOB_ENDED_NOTICE.format(
            namespace=NAMESPACE,
            job=job_to_run,
        ),
    )

    # Now fetch the logs!
    fetch_logs(logs_destination)


logging.basicConfig(format=LOG_FORMAT)

if __name__ == '__main__':
    arguments = docopt(
        __doc__,
        version='osp-pipeline Cluster Manager v%s' % __version__,
    )
    if arguments['deploy']:
        deploy_cluster(verbose=arguments['--very-verbose'])
    elif arguments['update']:
        update_cluster(verbose=arguments['--very-verbose'])
    elif arguments['run']:
        run_job(
            job_to_run=arguments['<job>'],
            logs_destination=arguments['<logs_destination>'],
            verbose=arguments['--very-verbose'],
        )
    elif arguments['logs']:
        fetch_logs(
            arguments['<destination>'],
            verbose=arguments['--very-verbose'],
        )
    elif arguments['masterui']:
        print('http://' + get_master_ip() + ':8080/')
    elif arguments['terminate']:
        terminate_spark_nodes(interactive=not arguments['--non-interactive'])
