"""Run the `extract_sales` ECS task."""

from datetime import datetime
from datetime import timezone
import json
import time

from lib import config
from lib.utils.aws import AwsCredentials
from lib.utils.aws import get_credentials_for_assumed_role
from lib.utils.aws import get_ec2_client
from lib.utils.aws import get_ecs_client
from lib.utils.aws import get_security_group
from lib.utils.aws import get_sts_client
from lib.utils.aws import get_subnet
from lib.utils.aws import get_vpc
from lib.utils.event import get_abacus_event

_ecs_credentials: AwsCredentials | None = None


def run_extract_sales_task(dag_run, **kwargs):
    """Run the `extract_sales` ECS task."""
    event = get_abacus_event(dag_run, **kwargs)
    batch_id = event.target_id

    task_instance = kwargs.get('task_instance')
    sales_ingest_type = task_instance.xcom_pull(
        task_ids='determine_sales_ingest_type',
        key='sales_ingest_type'
    )

    print(f'Running the extract_sales task for {sales_ingest_type}/{batch_id}')

    task_arn = _run_ecs_task(sales_ingest_type, batch_id)
    task_exit_code = _wait_for_ecs_task(task_arn)

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


def _get_ecs_credentials() -> AwsCredentials:
    """Return the ECS credentials if they are valid or generate them."""
    global _ecs_credentials

    if _are_ecs_credentials_valid():
        assert _ecs_credentials
        return _ecs_credentials

    print('Generating ECS credentials')

    _ecs_credentials = get_credentials_for_assumed_role(
        get_sts_client(),
        role_arn=config.EXTRACT_SALES_DAG_EXECUTION_ROLE,
        session_name='run_extract_sales_session'
    )

    return _ecs_credentials


def _are_ecs_credentials_valid() -> bool:
    """Check if the ECS credentials have been generated and have not expired."""
    global _ecs_credentials

    if _ecs_credentials is None:
        return False

    expiration = _ecs_credentials['Expiration']
    now = datetime.now(timezone.utc)
    duration_left = int(expiration.timestamp() - now.timestamp())

    # NOTE: If the credentials are scheduled to expire in the next minute
    # we consider them invalid and thus must be refreshed.
    if duration_left < 60:
        print(f'ECS credentials will expire in {duration_left}s')
        return False

    return True


def _get_network_config() -> dict:
    """Get the network configuration to run the ECS task."""
    credentials = _get_ecs_credentials()
    ec2_client = get_ec2_client(credentials)
    vpc_id = get_vpc(ec2_client)

    return {
        'awsvpcConfiguration': {
            'subnets': [get_subnet(ec2_client, vpc_id)],
            'securityGroups': [
                get_security_group(
                    ec2_client,
                    vpc_id,
                    config.EXTRACT_SALES_ECS_SERVICE_NAME
                )
            ],
            'assignPublicIp': 'DISABLED'
        }
    }


def _run_ecs_task(sales_ingest_type: str, batch_id: int):
    """Run the `extract_sales` ECS task.

    Args:
        sales_ingest_type (str): The type of sales to ingest.
        batch_id (int): The ID of the batch of sales to ingest.
    """
    credentials = _get_ecs_credentials()
    ecs_client = get_ecs_client(credentials)

    run_result = ecs_client.run_task(
        cluster=config.EXTRACT_SALES_ECS_SERVICE_NAME,
        taskDefinition=config.EXTRACT_SALES_ECS_SERVICE_NAME,
        launchType='FARGATE',
        networkConfiguration=_get_network_config(),
        overrides={
            'containerOverrides': [
                {
                    'name': 'ecs-abacus-extract-sales',
                    'environment': [
                        {
                            'name': 'SALES_TYPE',
                            'value': sales_ingest_type
                        },
                        {
                            'name': 'BATCH_ID',
                            'value': str(batch_id)
                        }
                    ]
                }
            ]
        }
    )

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


def _wait_for_ecs_task(task_arn) -> dict:
    """Wait for the ECS task to complete.

    Args:
        task_arn (str): The task ARN.

    Returns:
        dict: A dict with the exit code and reason.
    """
    print('Waiting for ECS task to complete')

    while True:
        credentials = _get_ecs_credentials()
        ecs_client = get_ecs_client(credentials)

        describe_result = ecs_client.describe_tasks(
            cluster=config.EXTRACT_SALES_ECS_SERVICE_NAME,
            tasks=[task_arn]
        )

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

        print(f'Task status: {task_status}')

        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(30)