"""pypyr step that prepares context for an ecs waiter."""
import logging
from pypyr.errors import KeyNotInContextError


# pypyr logger means the log level will be set correctly and output formatted.
logger = logging.getLogger(__name__)


def run_step(context):
    """Run me after an ecs task run or stop to prepare an ecs waiter.

    Prepares the awsWaitIn context key for pypyraws.steps.wait

    Use this step after any of the following methods if you want to use one of
    the ecs waiters to wait for a specific state:
    - describe_services
    - describe_tasks
    - list_services - specify awsEcsWaitPrepCluster if you don't want default
    - list_tasks - specify awsEcsWaitPrepCluster if you don't want default
    - run_task
    - start_task
    - stop_task
    - update_service

    You don't have to use this step, you could always just construct the
    awsWaitIn dictionary in context yourself. It just so happens this step
    saves you some legwork to do so.

    Args:
        context:
            Dictionary. Mandatory.
            Requires the following context keys in context:
                - awsClientOut. dict. mandatory. This is the context key that
                  any ecs command executed by pypyraws.steps.service adds.
                  Chances are pretty good you don't want to construct this by
                  hand yourself - the idea is to use the output as generated by
                  one of the supported ecs methods.
                - awsEcsWaitPrepCluster. string. optional. The short name or
                  full arn of the cluster that hosts the task to describe. If
                  you do not specify a cluster, the default cluster is assumed.
                  For most of the ecs methods the code automatically deduces
                  the cluster from awsClientOut, so don't worry about it.
                  But, when following list_services and list_tasks, you have to
                  specify this parameter. Specifying this parameter will
                  override any automatically deduced cluster arn.

    Returns:
        None.
        Overwrites the awsWaitIn key in context. The new awsWaitIn will contain
        waitArgs filled with the task or service arns found in awsClientOut.

    Raises:
        pypyr.errors.KeyNotInContextError: awsClientOut missing in context.
        pypyr.errors.KeyInContextHasNoValueError: awsClientOut exists but is
                                                  None.
    """
    logger.debug("started")

    context.assert_key_has_value('awsClientOut', __name__)

    if 'awsEcsWaitPrepCluster' in context:
        logger.debug("awsEcsWaitPrepCluster specified in input context")
        # awsEcsWaitPrepCluster always overrides automatically deduced cluster
        cluster = context['awsEcsWaitPrepCluster']
    else:
        cluster = None

    parse_me = context['awsClientOut']
    isTask = False
    isService = False

    if 'service' in parse_me:
        logger.debug("Found 'service' in awsClientOut")
        if cluster is None:
            cluster = parse_me['service']['clusterArn']
        arn_list = [parse_me['service']['serviceArn']]
        isService = True
    elif 'serviceArns' in parse_me:
        logger.debug("Found 'serviceArns' in awsClientOut")
        # no cluster arn is available for these
        arn_list = parse_me['serviceArns']
        isService = True
    elif 'services' in parse_me:
        logger.debug("Found 'services' in awsClientOut")
        if cluster is None:
            # so happens aws methods returning 'services' all do so with only 1
            # cluster input.
            cluster = parse_me['services'][0]['clusterArn']
        arn_list = [svc['serviceArn']
                    for svc in parse_me['services']]
        isService = True
    elif 'task' in parse_me:
        logger.debug("Found 'task' in awsClientOut")
        if cluster is None:
            cluster = parse_me['task']['clusterArn']
        arn_list = [parse_me['task']['taskArn']]
        isTask = True
    elif 'taskArns' in parse_me:
        logger.debug("Found 'taskArns' in awsClientOut")
        # no cluster arn is available for these
        arn_list = parse_me['taskArns']
        isTask = True
    elif 'tasks' in parse_me:
        logger.debug("Found 'tasks' in awsClientOut")
        if cluster is None:
            # so happens the aws methods that returns 'tasks' all do so with
            # only 1 cluster input, i.e cluster will be the same for all of the
            # return tasks
            cluster = parse_me['tasks'][0]['clusterArn']
        arn_list = [task['taskArn'] for task in parse_me['tasks']]
        isTask = True
    else:
        raise KeyNotInContextError('Run ecswaitprep after an ecs method that '
                                   'does something with services or tasks. '
                                   'Couldn\'t find service, serviceArns, '
                                   'services, task, taskArns or tasks in '
                                   'awsClientOut.')

    waiter_dict = {}
    if cluster is None:
        logger.debug("No cluster specified. Waiter will use default cluster.")
    else:
        logger.debug(f"{cluster} cluster specified.")
        waiter_dict['cluster'] = cluster

    if isTask:
        logger.debug("Adding task arns")
        waiter_dict['tasks'] = arn_list

    if isService:
        logger.debug("Adding service arns")
        waiter_dict['services'] = arn_list

    if 'awsWaitIn' not in context:
        context['awsWaitIn'] = {}
    context['awsWaitIn']['waitArgs'] = waiter_dict
    logger.info("added context['awsWaitIn']['waitArgs']")
    logger.debug("done")
