# Portions of this file are copied from the Yelp MRJob project:
#
#   https://github.com/Yelp/mrjob
#
#
# Copyright 2009-2013 Yelp, David Marin
# Copyright 2015 Yelp
# Copyright 2017 Yelp
# Copyright 2018 Contributors
# Copyright 2019 Yelp and Contributors
#
# 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 gzip
import re
import time
from io import BytesIO

import boto3
import six
from botocore.exceptions import WaiterError
from dagster_aws.utils.mrjob.utils import _boto3_now, _wrap_aws_client, strip_microseconds

import dagster
from dagster import check
from dagster.seven import urlparse

from .types import EMR_CLUSTER_TERMINATED_STATES, EmrClusterState, EmrStepState

# if we can't create or find our own service role, use the one
# created by the AWS console and CLI
_FALLBACK_SERVICE_ROLE = 'EMR_DefaultRole'

# if we can't create or find our own instance profile, use the one
# created by the AWS console and CLI
_FALLBACK_INSTANCE_PROFILE = 'EMR_EC2_DefaultRole'


class EmrError(Exception):
    pass


class EmrJobRunner:
    def __init__(
        self, region, check_cluster_every=30, aws_access_key_id=None, aws_secret_access_key=None,
    ):
        '''This object encapsulates various utilities for interacting with EMR clusters and invoking
        steps (jobs) on them.

        See also :py:class:`~dagster_aws.emr.EmrPySparkResource`, which wraps this job runner in a
        resource for pyspark workloads.

        Args:
            region (str): AWS region to use
            check_cluster_every (int, optional): How frequently to poll boto3 APIs for updates.
                Defaults to 30 seconds.
            aws_access_key_id ([type], optional): AWS access key ID. Defaults to None, which will
                use the default boto3 credentials chain.
            aws_secret_access_key ([type], optional): AWS secret access key. Defaults to None, which
                will use the default boto3 credentials chain.
        '''
        self.region = check.str_param(region, 'region')

        # This is in seconds
        self.check_cluster_every = check.int_param(check_cluster_every, 'check_cluster_every')
        self.aws_access_key_id = check.opt_str_param(aws_access_key_id, 'aws_access_key_id')
        self.aws_secret_access_key = check.opt_str_param(
            aws_secret_access_key, 'aws_secret_access_key'
        )

    def make_emr_client(self):
        '''Creates a boto3 EMR client. Construction is wrapped in retries in case client connection
        fails transiently.

        Returns:
            botocore.client.EMR: An EMR client
        '''
        raw_emr_client = boto3.client(
            'emr',
            aws_access_key_id=self.aws_access_key_id,
            aws_secret_access_key=self.aws_secret_access_key,
            region_name=self.region,
        )
        return _wrap_aws_client(raw_emr_client, min_backoff=self.check_cluster_every)

    def cluster_id_from_name(self, cluster_name):
        '''Get a cluster ID in the format "j-123ABC123ABC1" given a cluster name "my cool cluster".

        Args:
            cluster_name (str): The name of the cluster for which to find an ID

        Returns:
            str: The ID of the cluster

        Raises:
            EmrError: No cluster with the specified name exists
        '''
        check.str_param(cluster_name, 'cluster_name')

        response = self.make_emr_client().list_clusters().get('Clusters', [])
        for cluster in response:
            if cluster['Name'] == cluster_name:
                return cluster['Id']

        raise EmrError(
            'cluster {cluster_name} not found in region {region}'.format(
                cluster_name=cluster_name, region=self.region
            )
        )

    @staticmethod
    def construct_step_dict_for_command(step_name, command, action_on_failure='CONTINUE'):
        '''Construct an EMR step definition which uses command-runner.jar to execute a shell command
        on the EMR master.

        Args:
            step_name (str): The name of the EMR step (will show up in the EMR UI)
            command (str): The shell command to execute with command-runner.jar
            action_on_failure (str, optional): Configure action on failure (e.g., continue, or
                terminate the cluster). Defaults to 'CONTINUE'.

        Returns:
            dict: Step definition dict
        '''
        check.str_param(step_name, 'step_name')
        check.list_param(command, 'command', of_type=str)
        check.str_param(action_on_failure, 'action_on_failure')

        return {
            'Name': step_name,
            'ActionOnFailure': action_on_failure,
            'HadoopJarStep': {'Jar': 'command-runner.jar', 'Args': command},
        }

    def add_tags(self, log, tags, cluster_id):
        '''Add tags in the dict tags to cluster cluster_id.

        Args:
            log (DagsterLogManager): Log manager, for logging
            tags (dict): Dictionary of {'key': 'value'} tags
            cluster_id (str): The ID of the cluster to tag
        '''
        check.dict_param(tags, 'tags')
        check.str_param(cluster_id, 'cluster_id')

        tags_items = sorted(tags.items())

        self.make_emr_client().add_tags(
            ResourceId=cluster_id, Tags=[dict(Key=k, Value=v) for k, v in tags_items]
        )

        log.info(
            'Added EMR tags to cluster %s: %s'
            % (cluster_id, ', '.join('%s=%s' % (tag, value) for tag, value in tags_items))
        )

    def run_job_flow(self, log, cluster_config):
        '''Create an empty cluster on EMR, and return the ID of that job flow.

        Args:
            log (DagsterLogManager): Log manager, for logging
            cluster_config (dict): Configuration for this EMR job flow. See:
                https://docs.aws.amazon.com/emr/latest/APIReference/API_RunJobFlow.html

        Returns:
            str: The cluster ID, e.g. "j-ZKIY4CKQRX72"
        '''
        check.dict_param(cluster_config, 'cluster_config')

        log.debug('Creating Elastic MapReduce cluster')
        emr_client = self.make_emr_client()

        log.debug(
            'Calling run_job_flow(%s)'
            % (', '.join('%s=%r' % (k, v) for k, v in sorted(cluster_config.items())))
        )
        cluster_id = emr_client.run_job_flow(**cluster_config)['JobFlowId']

        log.info('Created new cluster %s' % cluster_id)

        # set EMR tags for the cluster
        tags = cluster_config.get('Tags', {})
        tags['__dagster_version'] = dagster.__version__
        self.add_tags(log, tags, cluster_id)
        return cluster_id

    def describe_cluster(self, cluster_id):
        '''Thin wrapper over boto3 describe_cluster.

        Args:
            cluster_id (str): Cluster to inspect

        Returns:
            dict: The cluster info. See:
                https://docs.aws.amazon.com/emr/latest/APIReference/API_DescribeCluster.html
        '''
        check.str_param(cluster_id, 'cluster_id')

        emr_client = self.make_emr_client()
        return emr_client.describe_cluster(ClusterId=cluster_id)

    def describe_step(self, cluster_id, step_id):
        '''Thin wrapper over boto3 describe_step.

        Args:
            cluster_id (str): Cluster to inspect
            step_id (str): Step ID to describe

        Returns:
            dict: The step info. See:
                https://docs.aws.amazon.com/emr/latest/APIReference/API_DescribeStep.html
        '''
        check.str_param(cluster_id, 'cluster_id')
        check.str_param(step_id, 'step_id')

        emr_client = self.make_emr_client()
        return emr_client.describe_step(ClusterId=cluster_id, StepId=step_id)

    def add_job_flow_steps(self, log, cluster_id, step_defs):
        '''Submit the constructed job flow steps to EMR for execution.

        Args:
            log (DagsterLogManager): Log manager, for logging
            cluster_id (str): The ID of the cluster
            step_defs (List[dict]): List of steps; see also `construct_step_dict_for_command`

        Returns:
            List[str]: list of step IDs.
        '''
        check.str_param(cluster_id, 'cluster_id')
        check.list_param(step_defs, 'step_defs', of_type=dict)

        emr_client = self.make_emr_client()

        steps_kwargs = dict(JobFlowId=cluster_id, Steps=step_defs)
        log.debug(
            'Calling add_job_flow_steps(%s)'
            % ','.join(('%s=%r' % (k, v)) for k, v in steps_kwargs.items())
        )
        return emr_client.add_job_flow_steps(**steps_kwargs)['StepIds']

    def wait_for_emr_steps_to_complete(self, log, cluster_id, emr_step_ids):
        '''Wait for every step of the job to complete, one by one.

        Args:
            log (DagsterLogManager): Log manager, for logging
            cluster_id (str): The ID of the cluster
            step_ids (List[str]): List of EMR step IDs to wait for
        '''
        check.str_param(cluster_id, 'cluster_id')
        check.list_param(emr_step_ids, 'step_ids', of_type=str)

        for emr_step_id in emr_step_ids:
            log.info('Waiting for EMR step %s to complete...' % emr_step_id)
            self._wait_for_emr_step_to_complete(log, cluster_id, emr_step_id)

    def _wait_for_emr_step_to_complete(self, log, cluster_id, emr_step_id):
        '''Helper for wait_for_steps_to_complete(). Wait for step with the given ID to complete.

        Args:
            cluster_id (str): The ID of the cluster
            emr_step_id (str): EMR Step ID to wait for

        Raises:
            EmrError: Raised when the step is marked by EMR as failed instead of completing
                successfully.
        '''
        check.str_param(cluster_id, 'cluster_id')
        check.str_param(emr_step_id, 'emr_step_id')

        while True:
            # don't antagonize EMR's throttling
            log.debug('Waiting %.1f seconds...' % self.check_cluster_every)
            time.sleep(self.check_cluster_every)

            step = self.describe_step(cluster_id, emr_step_id)['Step']
            step_state = EmrStepState(step['Status']['State'])

            if step_state == EmrStepState.Pending:
                cluster = self.describe_cluster(cluster_id)['Cluster']

                reason = _get_reason(cluster)
                reason_desc = (': %s' % reason) if reason else ''

                log.info('PENDING (cluster is %s%s)' % (cluster['Status']['State'], reason_desc))
                continue

            elif step_state == EmrStepState.Running:
                time_running_desc = ''

                start = step['Status']['Timeline'].get('StartDateTime')
                if start:
                    time_running_desc = ' for %s' % strip_microseconds(_boto3_now() - start)

                log.info('RUNNING%s' % time_running_desc)
                continue

            # we're done, will return at the end of this
            elif step_state == EmrStepState.Completed:
                log.info('COMPLETED')
                return
            else:
                # step has failed somehow. *reason* seems to only be set
                # when job is cancelled (e.g. 'Job terminated')
                reason = _get_reason(step)
                reason_desc = (' (%s)' % reason) if reason else ''

                log.info('%s%s' % (step_state.value, reason_desc))

                # print cluster status; this might give more context
                # why step didn't succeed
                cluster = self.describe_cluster(cluster_id)['Cluster']
                reason = _get_reason(cluster)
                reason_desc = (': %s' % reason) if reason else ''
                log.info(
                    'Cluster %s %s %s%s'
                    % (
                        cluster['Id'],
                        'was' if 'ED' in cluster['Status']['State'] else 'is',
                        cluster['Status']['State'],
                        reason_desc,
                    )
                )

                if EmrClusterState(cluster['Status']['State']) in EMR_CLUSTER_TERMINATED_STATES:
                    # was it caused by IAM roles?
                    self._check_for_missing_default_iam_roles(log, cluster)

                    # TODO: extract logs here to surface failure reason
                    # See: https://github.com/dagster-io/dagster/issues/1954

            if step_state == EmrStepState.Failed:
                log.info('EMR step %s failed' % emr_step_id)

            raise EmrError('EMR step failed')

    def _check_for_missing_default_iam_roles(self, log, cluster):
        '''If cluster couldn't start due to missing IAM roles, tell user what to do.'''

        check.dict_param(cluster, 'cluster')

        reason = _get_reason(cluster)
        if any(
            reason.endswith('/%s is invalid' % role)
            for role in (_FALLBACK_INSTANCE_PROFILE, _FALLBACK_SERVICE_ROLE)
        ):
            log.warning(
                'IAM roles are missing. See documentation for IAM roles on EMR here: '
                'https://docs.aws.amazon.com/emr/latest/ManagementGuide/emr-iam-roles.html'
            )

    def log_location_for_cluster(self, cluster_id):
        '''EMR clusters are typically launched with S3 logging configured. This method inspects a
        cluster using boto3 describe_cluster to retrieve the log URI.

        Args:
            cluster_id (str): The cluster to inspect.

        Raises:
            EmrError: the log URI was missing (S3 log mirroring not enabled for this cluster)

        Returns:
            (str, str): log bucket and key
        '''
        check.str_param(cluster_id, 'cluster_id')

        # The S3 log URI is specified per job flow (cluster)
        log_uri = self.describe_cluster(cluster_id)['Cluster'].get('LogUri', None)

        # ugh, seriously boto3?! This will come back as string "None"
        if log_uri == 'None' or log_uri is None:
            raise EmrError('Log URI not specified, cannot retrieve step execution logs')

        # For some reason the API returns an s3n:// protocol log URI instead of s3://
        log_uri = re.sub('^s3n', 's3', log_uri)
        log_uri_parsed = urlparse(log_uri)
        log_bucket = log_uri_parsed.netloc
        log_key_prefix = log_uri_parsed.path.lstrip('/')
        return log_bucket, log_key_prefix

    def retrieve_logs_for_step_id(self, log, cluster_id, step_id):
        '''Retrieves stdout and stderr logs for the given step ID.

        Args:
            log (DagsterLogManager): Log manager, for logging
            cluster_id (str): EMR cluster ID
            step_id (str): EMR step ID for the job that was submitted.

        Returns
            (str, str): Tuple of stdout log string contents, and stderr log string contents
        '''
        check.str_param(cluster_id, 'cluster_id')
        check.str_param(step_id, 'step_id')

        log_bucket, log_key_prefix = self.log_location_for_cluster(cluster_id)

        prefix = '{log_key_prefix}{cluster_id}/steps/{step_id}'.format(
            log_key_prefix=log_key_prefix, cluster_id=cluster_id, step_id=step_id
        )
        stdout_log = self.wait_for_log(log, log_bucket, '{prefix}/stdout.gz'.format(prefix=prefix))
        stderr_log = self.wait_for_log(log, log_bucket, '{prefix}/stderr.gz'.format(prefix=prefix))
        return stdout_log, stderr_log

    def wait_for_log(self, log, log_bucket, log_key, waiter_delay=30, waiter_max_attempts=20):
        '''Wait for gzipped EMR logs to appear on S3. Note that EMR syncs logs to S3 every 5
        minutes, so this may take a long time.

        Args:
            log_bucket (str): S3 bucket where log is expected to appear
            log_key (str): S3 key for the log file
            waiter_delay (int): How long to wait between attempts to check S3 for the log file
            waiter_max_attempts (int): Number of attempts before giving up on waiting

        Raises:
            EmrError: Raised if we waited the full duration and the logs did not appear

        Returns:
            str: contents of the log file
        '''
        check.str_param(log_bucket, 'log_bucket')
        check.str_param(log_key, 'log_key')
        check.int_param(waiter_delay, 'waiter_delay')
        check.int_param(waiter_max_attempts, 'waiter_max_attempts')

        log.info(
            'Attempting to get log: s3://{log_bucket}/{log_key}'.format(
                log_bucket=log_bucket, log_key=log_key
            )
        )

        s3 = _wrap_aws_client(boto3.client('s3'), min_backoff=self.check_cluster_every)
        waiter = s3.get_waiter('object_exists')
        try:
            waiter.wait(
                Bucket=log_bucket,
                Key=log_key,
                WaiterConfig={'Delay': waiter_delay, 'MaxAttempts': waiter_max_attempts},
            )
        except WaiterError as err:
            six.raise_from(
                EmrError('EMR log file did not appear on S3 after waiting'), err,
            )
        obj = BytesIO(s3.get_object(Bucket=log_bucket, Key=log_key)['Body'].read())
        gzip_file = gzip.GzipFile(fileobj=obj)
        return gzip_file.read().decode('utf-8')


def _get_reason(cluster_or_step):
    '''Get state change reason message.'''
    # StateChangeReason is {} before the first state change
    return cluster_or_step['Status']['StateChangeReason'].get('Message', '')
