"""Utility shared functions."""

import datetime
import time
from typing import Any

from fargate_tools.cli.constants import VerifyMode


def poll_task_status(
    client: Any,
    verify_mode: VerifyMode,
    update_timeout: int,
    cluster_name: str,
    service_name: str,
    task_definition_arn: str,
    container_name: str,
    grace_period: int,
) -> bool:
    """Poll service to see if task has launched."""
    timeout = time.time() + update_timeout
    print(f'\nPolling for new task status. Timeout: {update_timeout} seconds')
    while time.time() < timeout:
        print(datetime.datetime.now())
        paginator = client.get_paginator('list_tasks')
        page_iterator = paginator.paginate(
            cluster=cluster_name, serviceName=service_name)

        # Get all tasks and look for our new one.
        for task_response in page_iterator:
            task_arns = task_response['taskArns']
            if not task_arns:
                continue
            tasks = client.describe_tasks(
                cluster=cluster_name, tasks=task_arns)['tasks']
            for task_description in tasks:
                if task_definition_arn in task_description['taskDefinitionArn']:
                    # This is our new task. Focus on it specifically.
                    task_arn = task_description['taskArn']
                    print(f'\nTask found. Status: {task_description["lastStatus"]}')
                    return _verify_task(
                        client=client,
                        verify_mode=verify_mode,
                        cluster_name=cluster_name,
                        task_arn=task_arn,
                        container_name=container_name,
                        timeout=timeout,
                        grace_period=grace_period
                    )
        time.sleep(15)
    print('\nTimed out waiting for new task to appear. Deployment failed.')
    return False


def _verify_task(
    client: Any,
    verify_mode: VerifyMode,
    cluster_name: str,
    task_arn: str,
    container_name: str,
    timeout: float,
    grace_period: int,
) -> bool:
    print(f'Verifying task {task_arn}')
    print(f'Verify mode is {verify_mode}')
    if verify_mode == VerifyMode.HEALTH_CHECK:
        return _verify_task_healthy(
            client=client,
            cluster_name=cluster_name,
            task_arn=task_arn,
            timeout=timeout,
            grace_period=grace_period
        )
    elif verify_mode == VerifyMode.EXIT_CODE:
        return _verify_task_exit_code(
            client=client,
            cluster_name=cluster_name,
            task_arn=task_arn,
            container_name=container_name,
            timeout=timeout,
        )
    elif verify_mode == VerifyMode.TASK_RUNNING:
        return _verify_task_running(
            client=client,
            cluster_name=cluster_name,
            task_arn=task_arn,
            timeout=timeout
        )
    elif verify_mode == VerifyMode.OFF:
        return True
    else:
        raise Exception(f'Unsupported verify_mode {verify_mode}.')


def _verify_task_healthy(
    client: Any,
    cluster_name: str,
    task_arn: str,
    timeout: float,
    grace_period: int,
) -> bool:
    health_check_grace_period = None
    while time.time() < timeout:
        for new_task in client.describe_tasks(
                cluster=cluster_name, tasks=[task_arn])['tasks']:
            if new_task['lastStatus'] == 'RUNNING' and not health_check_grace_period:
                """
                Once task is running, start health check
                grace period. Tasks marked as unhealthy should
                not be killed during this time.
                """
                health_check_grace_period = time.time() + grace_period
                print('\nBegin health check grace period')
            if new_task['healthStatus'] == 'HEALTHY':
                print(f'\nTask status: {new_task["lastStatus"]}')
                print(f'Task health: {new_task["healthStatus"]}')
                print('Task is healthy. Deployment complete.')
                return True
            elif new_task['healthStatus'] == 'UNHEALTHY':
                if time.time() < health_check_grace_period:
                    print('\nTask is unhealthy. Waiting for grace period')
                    print(datetime.datetime.now())
                    time.sleep(5)
                else:
                    print('Task is unhealthy. Deployment failed.')
                    return False
            elif new_task['lastStatus'] == 'STOPPED':
                print('Task was stopped. Deployment failed.')
                return False
            else:
                print(f'\nTask status: {new_task["lastStatus"]}')
                print(f'Task health: {new_task["healthStatus"]}')
                print(datetime.datetime.now())
                time.sleep(5)
    print('\nTimed out waiting for task to become healthy. Deployment failed.')
    return False


def _verify_task_exit_code(
    client: Any,
    cluster_name: str,
    task_arn: str,
    container_name: str,
    timeout: float,
) -> bool:
    while time.time() < timeout:
        for new_task in client.describe_tasks(
                cluster=cluster_name, tasks=[task_arn])['tasks']:
            if 'stopCode' in new_task:
                stop_code = new_task['stopCode']
                if stop_code != 'EssentialContainerExited':
                    print(f'\nTask stopped with unexpected stop code {stop_code}. Deployment failed.')
                    return False
                container = next(
                    container for container in new_task['containers']
                    if container['name'] == container_name
                )
                container_exit_code = container['exitCode']
                if container_exit_code != 0:
                    print(f'\nContainer {container["name"]} exited with code {container_exit_code}. Deployment failed.')
                    return False
                else:
                    print(f'\nContainer {container["name"]} exited with code {container_exit_code}. Deployment complete.')
                    return True
            else:
                print(f'\nTask status: {new_task["lastStatus"]}')
                print(datetime.datetime.now())
                time.sleep(5)
    print('\nTimed out waiting for task to exit with code 0. Deployment failed.')
    return False


def _verify_task_running(
    client: Any,
    cluster_name: str,
    task_arn: str,
    timeout: float,
) -> bool:
    while time.time() < timeout:
        for new_task in client.describe_tasks(
                cluster=cluster_name, tasks=[task_arn])['tasks']:
            last_status = new_task['lastStatus']
            if last_status == 'RUNNING':
                print('Task is running. Deployment complete.')
                return True
            else:
                print(f'\nTask status: {new_task["lastStatus"]}')
                print(datetime.datetime.now())
                if last_status == 'STOPPED':
                    return False
                time.sleep(5)
    return False