import boto.emr
from boto.emr.instance_group import InstanceGroup
from boto.emr.step import HiveStep
from boto.emr.step import JarStep
from garcon import task


# Default EMR settings
STEPS = [
    JarStep(
        name='Setup hadoop',
        jar='s3://elasticmapreduce/libs/script-runner/script-runner.jar',
        step_args=['s3://elasticmapreduce/libs/state-pusher/0.1/fetch'])]

INSTANCE_GROUPS = [
    InstanceGroup(
        1, 'MASTER', 'm1.medium', 'ON_DEMAND', 'Master instance group'),
    InstanceGroup(
        2, 'CORE', 'm1.medium', 'ON_DEMAND', 'Core instance group')]


@task.decorate(timeout=7000)
def add_hive_settings_step(
        activity, emr_cluster_id, step_name, step_args):
    """Modify hive settings

    Args:
        activity (ActivityWorker): the swf activity worker.
        emr_cluster_id (str): the emr cluster id.
        step_name (str): name of the hive step
        step_args (list): list of arguments for the hive step
    """

    connector = boto.emr.connect_to_region('us-east-1')
    response = connector.add_jobflow_steps(
        emr_cluster_id, [
            JarStep(
                name=step_name,
                jar=(
                    's3://elasticmapreduce/libs/script-runner/'
                    'script-runner.jar'),
                step_args=step_args)
        ])
    step_id = response.stepids[0].value

    return {'emr.step_id': step_id}


@task.decorate(timeout=7200)
def launch_emr_cluster(
        activity, emr_cluster_id, path_to_logs, job_name, tag_name,
        ec2_instance_name, keep_alive, instance_groups):
    """Elastic Map Reduce for generating reports from varieties of data
       dependencies

    Args:
        activity (ActivityWorker): The activity worker.
        emr_cluster_id (int): the cluster id.
        path_to_logs (str): destination of the logs.
        job_name (str): name of the job.
        tag_name (str): the tag to add to the emr cluster.
        ec2_instance_name (str): optional, name of the ec2 instances that will
        be used by the cluster.
        keep_alive (boolean): indicate whether cluster should shutdown
        immediately after job is done.
        instance_groups (optional[list]): list of InstanceGroup settings

    Return:
        dict: informations about the cluster (mostly its id.)
    """

    if emr_cluster_id:
        activity.logger.info(
            'Launch of the emr cluster: found a emr cluster already available '
            'to execute the step (cluster id: {cluster_id})'.format(
                cluster_id=emr_cluster_id))
        return {'emr.cluster_id': emr_cluster_id}

    activity.logger.info('Launch the emr cluster')
    connector = boto.emr.connect_to_region('us-east-1')

    if keep_alive is None:
        keep_alive = True
    # ec2_keyname='vector-dev-AWS-US-East',

    if instance_groups is None:
        instance_groups = INSTANCE_GROUPS

    api_params = {"Instances.Ec2SubnetId": "subnet-d63d81a0"}

    cluster_id = connector.run_jobflow(
        ami_version='3.11.0',
        api_params=api_params,
        action_on_failure='CONTINUE',
        ec2_keyname='dev_orchard_admin',
        instance_groups=instance_groups,
        job_flow_role='EMR_EC2_DefaultRole',
        keep_alive=keep_alive,
        log_uri=path_to_logs,
        name=job_name.upper(),
        visible_to_all_users=True,
        steps=STEPS,
        service_role='EMR_DefaultRole')

    tags = {}

    if ec2_instance_name:
        activity.logger.info(
            'Add EC2 instance name to the emr cluster ec2 instances')
        tags.update(Name=ec2_instance_name)

    if tag_name:
        activity.logger.info(
            'Add tag name ({tag_name}) to the EMR cluster'.format(
                tag_name=tag_name))
        tags.update({tag_name: None})

    if tags:
        connector.add_tags(cluster_id, tags)

    return {'emr.cluster_id': cluster_id}


@task.decorate(timeout=140)
def add_streaming_job_step(
        activity, emr_cluster_id, step_name, files, mapper, reducer, input,
        output, jobconfs):
    """Add a streaming job step to an EMR cluster

    Args:
        activity (ActivityWorker): the swf activity worker.
        emr_cluster_id (str): the emr cluster id.
        step_name (str): name of the step.
        files (list): list of the full path to the mapper and reducer files
        on s3.
        mapper (str): mapper s3 file name.
        reducer (str): reducer s3 file name.
        input (str): full path to the source data file(s).
        output (str): path to the directory fo the destination.
        jobconfs (list): list of string of jobconf argument
    """

    activity.logger.info(
        'Add streaming job step to cluster (id: {cluster_id})'.format(
            cluster_id=emr_cluster_id))

    connector = boto.emr.connect_to_region('us-east-1')
    step_args = [
        '-files',
        ','.join(files),
        '-mapper', mapper,
        '-reducer', reducer,
        '-input', input,
        '-output', output
    ]

    if len(jobconfs) > 0:
        for jobconf in jobconfs:
            step_args.append('-jobconf')
            step_args.append(jobconf)

    response = connector.add_jobflow_steps(
        emr_cluster_id, [
            JarStep(
                name=step_name,
                jar='/home/hadoop/contrib/streaming/hadoop-streaming.jar',
                step_args=step_args)])

    step_id = response.stepids[0].value
    activity.logger.info(
        'Emr report generation step successfully added (id: {step_id})'.format(
            step_id=step_id))

    return {'emr.step_id': step_id}


@task.decorate(timeout=7000)
def add_s3_distcp_step(
        activity, emr_cluster_id, source_path, destination_path):
    """Copy files from HDFS to S3

    Args:
        activity (ActivityWorker): the swf activity worker.
        emr_cluster_id (str): the emr cluster id.
        source_path (str): HDFS or S3 path. Example: hdfs:///dev/test
        destination_path (str): HDFS or S3 path. Example: s3://bucket/test/
    """

    connector = boto.emr.connect_to_region('us-east-1')
    response = connector.add_jobflow_steps(
        emr_cluster_id, [
            JarStep(
                name='Transport Files Between HDFS and S3',
                jar='/home/hadoop/lib/emr-s3distcp-1.0.jar',
                step_args=[
                    '--src',
                    source_path,
                    '--dest',
                    destination_path
                ])
        ])
    step_id = response.stepids[0].value
    activity.logger.info('Files have been copied from HDFS to S3')

    return {'emr.step_id': step_id}


@task.decorate(timeout=2000)
def add_hive_job_step(
        activity, emr_cluster_id, source_s3_path, step_name,
        destination_s3_path, hql_s3_path):
    """Add a step to an EMR cluster.

    Args:
        activity (ActivityWorker): the swf activity worker.
        emr_cluster_id (str): the emr cluster id.
        source_s3_path (str): input data (s3 path.)
        step_name (str): name of each EMR step
        destination_s3_path (str): s3 path stores report
        hql_s3_path (str): s3 path stores hql

    Return:
        dict: contains the emr step id.
    """

    activity.logger.info(
        'Add step to report generation cluster (id: {cluster_id})'.format(
            cluster_id=emr_cluster_id))

    connector = boto.emr.connect_to_region('us-east-1')
    hive_args = []

    if source_s3_path:
        hive_args.append('-d')
        hive_args.append('INPUT={}'.format(source_s3_path))

    if destination_s3_path:
        hive_args.append('-d')
        hive_args.append('OUTPUT={}'.format(destination_s3_path))

    response = connector.add_jobflow_steps(
        emr_cluster_id, [
            HiveStep(
                name='Hive Program ({})'.format(step_name),
                hive_file=hql_s3_path,
                hive_versions='0.13.1',
                hive_args=hive_args)])

    step_id = response.stepids[0].value
    activity.logger.info(
        'Emr report generation step successfully added (id: {step_id})'.format(
            step_id=step_id))

    return {'emr.step_id': step_id}