import json
import time

import boto3

from lib.config import EXTRACT_SALES_SERVICE_NAME
from lib.config import SUBNET_IDS


def run_extract_sales_task(**kwargs):
    batch_id = kwargs.get('params').get('batch_id')

    ecs_client = _get_ecs_client()
    task_arn = _run_ecs_task(ecs_client, batch_id)
    task_exit_code = _wait_for_ecs_task(ecs_client, task_arn)

    if task_exit_code['exit_code'] != 0:
        raise Exception(
            f'Error running task: {json.dumps(task_exit_code)}'
        )


def _get_ecs_client():
    return boto3.client(
        'ecs',
        aws_access_key_id='localstack',
        aws_secret_access_key='localstack',
        region_name='us-east-1',
        endpoint_url='http://localstack:4566'
    )


def _run_ecs_task(ecs_client, batch_id):
    run_result = ecs_client.run_task(
        cluster=EXTRACT_SALES_SERVICE_NAME,
        taskDefinition=EXTRACT_SALES_SERVICE_NAME,
        launchType='FARGATE',
        networkConfiguration={
            'awsvpcConfiguration': {
                'subnets': SUBNET_IDS.split(','),
                'assignPublicIp': 'DISABLED'
            }
        },
        overrides={
            'containerOverrides': [
                {
                    'environment': [
                        {
                            'name': 'BATCH_ID',
                            'value': batch_id
                        }
                    ]
                }
            ]
        }
    )

    task_arn = run_result['tasks'][0]['taskArn']
    return task_arn


# TODO: Use: https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/ecs/waiter/TasksStopped.html
def _wait_for_ecs_task(ecs_client, task_arn):
    while True:
        describe_result = ecs_client.describe_tasks(
            cluster=EXTRACT_SALES_SERVICE_NAME,
            tasks=[task_arn]
        )

        task = describe_result['tasks'][0]
        task_status = task['lastStatus']

        if task_status == 'STOPPED':
            exit_code = task['containers'][0].get('exitCode', 0)
            exit_reason = task['containers'][0].get('reason')
            stop_code = task.get('stopCode')
            stopped_reason = task.get('stoppedReason')

            return {
                'exit_code': exit_code,
                'exit_reason': exit_reason,
                'stop_code': stop_code,
                'stopped_reason': stopped_reason
            }

        time.sleep(5)