from time import sleep

import boto3

from garcon import task

from garcon_contrib.aws.utils import garcon_emr_cluster


@task.decorate(timeout=1400)
def get_emr_cluster_by_tag(activity, tag_name):
    """Get an emr cluster with a tag name.

    Args:
        activity (ActivityWorker): The swf activity worker.
        tag_name (str): the tag name.
    """

    activity.logger.info(
        'Get the emr cluster by using {tag_name}'.format(
            tag_name=tag_name))

    cluster_id = garcon_emr_cluster.find_active_emr_cluster_id_by_tag(tag_name)
    if cluster_id:
        activity.logger.info(
            'Getting the emr cluster by using {tag_name} was successful '
            '(cluster id: {cluster_id})'.format(
                cluster_id=cluster_id,
                tag_name=tag_name))
        return {'emr.cluster_id': cluster_id}


@task.decorate(timeout=14000)
def completed_emr_job(activity, emr_cluster_id):
    """Task that check if the EMR cluster has completed all its steps.

    This task look for the last step every 5 minutes, if this step has
    completed, the task exit (which mark it as completed.) Failures or early
    termination of the step will end up marking the activity as failed.

    Args:
        activity (ActivityWorker): The swf activity worker.
        emr_cluster_id (str): the cluster id.
    """

    while True:
        sleep(300)

        emr = boto3.client('emr', region_name='us-east-1')
        steps = emr.list_steps(ClusterId=emr_cluster_id)['Steps']
        last_step = steps[-1]

        if last_step['Status']['State'] == 'COMPLETED':
            activity.logger.info(
                'The EMR cluster has completed all the jobs (cluster id: '
                '{cluster_id})'.format(
                    cluster_id=emr_cluster_id))
            return

        elif last_step['Status']['State'] in ['FAILED', 'TERMINATED']:
            activity.logger.info(
                'The EMR cluster has failed or has been terminated (cluster '
                'id: {cluster_id})'.format(
                    cluster_id=emr_cluster_id))

            activity.fail(
                reason=(
                    'Step %(name)s of the EMR job has either failed or had '
                    'been terminated.') % dict(
                        name=last_step['Name']))
            return

        activity.heartbeat(details=(
            'EMR Cluster is running the step "%(name)s"') % dict(
                name=last_step['Name']))


@task.decorate()
def shutdown_waiting_emr_cluster(activity, emr_cluster_id):
    """Shutdown the EMR cluster.

    Args:
        activity (ActivityWorker): The swf activity worker.
        emr_cluster_id (str): the cluster id.
    """

    emr = boto3.client('emr', region_name='us-east-1')
    cluster_details = emr.describe_cluster(ClusterId=emr_cluster_id)['Cluster']
    if cluster_details['Status']['State'] == 'WAITING':
        activity.logger.info(
            'Shutdown the emr cluster {emr_cluster_id}'.format(
                emr_cluster_id=emr_cluster_id))
        emr.terminate_job_flows(JobFlowIds=[emr_cluster_id])
    else:
        activity.logger.info(
            'Emr cluster {emr_cluster_id} is busy.'
            'The task will not shut it down'.format(
                emr_cluster_id=emr_cluster_id))


@task.decorate(timeout=600)
def is_emr_ready(activity, emr_cluster_id):
    """Check if an EMR cluster is ready.

    If you have a dependency on an EMR cluster: you will want to use this task.
    It makes sure the EMR cluster is up and running before your application
    start registering steps that have other dependencies.

    Args:
        activity (ActivityWorker): the swf activity worker.
        emr_cluster_id (str): the emr cluster id.
    """

    emr = boto3.client('emr', region_name='us-east-1')

    while True:
        state = emr.describe_cluster(
            ClusterId=emr_cluster_id)['Cluster']['Status']['State']
        activity.logger.info(
            'Wait for the emr cluster {cluster_id} to be ready (current '
            'status: {state})'.format(
                cluster_id=emr_cluster_id,
                state=state))

        if state in ['RUNNING', 'WAITING']:
            activity.logger.info('Emr cluster {cluster_id} is ready.'.format(
                cluster_id=emr_cluster_id))
            return

        if state in ['STARTING']:
            activity.heartbeat(details='EMR Cluster is starting')

        sleep(10)  # Sleep for 10 seconds


@task.decorate(7200)
def wait_emr_step_completed(activity, emr_cluster_id, emr_step_id):
    """Wait for an emr step to have completed.

    Args:
        activity (ActivityWorker): The swf activity worker.
        emr_cluster_id (str): the emr cluster id.
        emr_step_id (str): the emr step id the task will be waiting on.
    """

    emr = boto3.client('emr', region_name='us-east-1')

    while True:
        step = emr.describe_step(
            ClusterId=emr_cluster_id, StepId=emr_step_id)['Step']
        if not step:
            raise Exception(
                'The step could not be found for this flow. Was it created?')
        if step['Status']['State'] in ['CANCELLED', 'FAILED']:
            raise Exception(
                'The step has heen cancelled or has failed.')
        elif step['Status']['State'] in ['COMPLETED']:
            return

        activity.logger.info(
            'Wait for the emr step to complete (step id: {step_id}).'.format(
                step_id=emr_step_id))
        activity.heartbeat(details='EMR Job is in progress.')
        sleep(30)