"""Small helper script to run a Fargate task for swf workflows."""
import argparse
import logging
import os
import time

import boto3
from botocore.exceptions import ClientError
from dotenv import load_dotenv
from environs import Env

from fargate_tools.cli.constants import VerifyMode
from fargate_tools.cli.utils import _verify_task

# Use env for parsing environment variables
env = Env()

# Load env file if it exists
load_dotenv(verbose=True)

logging.basicConfig(level=logging.INFO)

ENVIRONMENT = os.environ.get('Environment', 'dev')
AWS_REGION = os.environ.get('AWS_REGION', 'us-east-1')
CLUSTER_NAME = os.environ.get('CLUSTER_NAME')
NUM_TASKS_TO_RUN = int(os.environ.get('NUM_TASKS_TO_RUN', 1))
SERVICE_NAME = os.environ.get('SERVICE_NAME')
CONTAINER_NAME = os.environ.get('CONTAINER_NAME', SERVICE_NAME)
SKIP_HEALTHCHECK = os.environ.get('SKIP_HEALTHCHECK')
TASK_GRACE_PERIOD = int(os.environ.get('TASK_GRACE_PERIOD', 120))
TASK_POLLING_TIMEOUT = int(os.environ.get('TASK_POLLING_TIMEOUT', 300))
TASK_STARTED_BY = os.environ.get('TASK_STARTED_BY', 'jenkins')

VPC_ENVIRONMENT_MAPPING = {
    'dev': 'dev',
    'qa': 'prod',
    'uat': 'prod',
    'prod': 'prod',
}

VPC_NAME = os.environ.get('VPC_NAME', VPC_ENVIRONMENT_MAPPING[ENVIRONMENT])

# see constants.py for details
verify_mode = env.enum(
    "VERIFY_MODE", enum=VerifyMode, default=VerifyMode.HEALTH_CHECK._name_)


def get_vpc_id(client):
    """Get VPC corresponding to environment

    Args:
        client (boto3.client): EC2 client
    Raises:
        SystemExit: If VPC cannot be found

    Returns:
        str: ID of VPC
    """
    response = client.describe_vpcs(
        Filters=[
            {
                'Name': 'tag:Name',
                'Values': [
                    VPC_NAME,
                ]
            },
        ],
    )
    if response['Vpcs']:
        vpc_id = response['Vpcs'][0]['VpcId']
        return vpc_id
    else:
        logging.info('VPC named {} not found'.format(
            VPC_ENVIRONMENT_MAPPING[ENVIRONMENT]))
        raise SystemExit('Could not retrieve VPC ID')


def get_security_group(client, vpc_id):
    """Returns security group corresponding to service

    Args:
        client (boto3.client): EC2 client
        vpc_id (string): ID of VPC for use in lookup filter

    Raises:
        SystemExit: If security group cannot be found

    Returns:
        str: ID of security group for service
    """
    security_group_name = '{}-{}-task-security-group'.format(
        ENVIRONMENT, SERVICE_NAME)
    response = client.describe_security_groups(
        Filters=[
            {
                'Name': 'group-name',
                'Values': [
                    security_group_name,
                ]
            },
            {
                'Name': 'vpc-id',
                'Values': [
                    vpc_id,
                ]
            },
        ],
    )

    if response['SecurityGroups']:
        """
        Infrastructure automation enforces uniqueness, but
        just in case, pick the first returned group
        """
        security_group_id = response['SecurityGroups'][0]['GroupId']
        return security_group_id
    else:
        logging.info('Security group named {} not found'.format(
            security_group_name))
        raise SystemExit('Could not retrieve security group')


def get_subnets(client, vpc_id):
    """Returns subnet IDs with most available IP addresses

    Args:
        client (boto3.client): EC2 client
        vpc_id (string): ID of VPC for use in lookup filter

    Raises:
        SystemExit: If subnets cannot be found

    Returns:
        list: IDs of subnets
    """
    response = client.describe_subnets(
        Filters=[
            {
                'Name': 'tag:Name',
                'Values': [
                    '*private*',
                ]
            },
            {
                'Name': 'tag:tier',
                'Values': [
                    'private',
                ]
            },
            {
                'Name': 'vpc-id',
                'Values': [
                    vpc_id,
                ]
            },
        ],
    )
    if response['Subnets']:
        """
        Find and return up to two subnets with the most available IP addresses
        """
        # Filter out subnets with less than 25 available IPs
        available_subnets = [subnet for subnet in response['Subnets'] if subnet['AvailableIpAddressCount'] > 25]

        # Filter subnets in use1-az3 zone
        eligible_availability_zones = [subnet for subnet in available_subnets if subnet['AvailabilityZoneId'] != 'use1-az3']

        sorted_subnets = sorted(
            eligible_availability_zones,
            key=lambda k: k['AvailableIpAddressCount'],
            reverse=True
        )
        subnet_ids = [subnet['SubnetId'] for subnet in sorted_subnets]
        return subnet_ids[0:2]
    else:
        logging.info('No subnets found in VPC ID {}'.format(vpc_id))
        raise SystemExit('Could not retrieve subnet IDs')


def run_task(client, security_group_id, subnet_ids, environment_overrides, commands_list):
    """Runs task and returns info about it, once started

    Args:
        client (boto3.client): ECS client
        security_group_id (str): security group used in task network config
        subnet_ids (list): subnets used in task network config
        environment_overrides (list): a list of environment variable overrides
        commands_list (list(str)): a list of command arguments for the task to run

    Raises:
        SystemExit: if an error occurs while running task

    Returns:
        str: ARN of newly started task
    """
    service_overrides = {}

    if environment_overrides:
        service_overrides['environment'] = environment_overrides
    if commands_list:
        service_overrides['command'] = commands_list

    if service_overrides:
        service_overrides['name'] = SERVICE_NAME
        container_overrides = [
            service_overrides
        ]
        logging.info(f'overrides are {container_overrides}')
    else:
        container_overrides = []
        logging.info(f'overrides are empty: {container_overrides}')

    try:
        response = client.run_task(
            cluster=CLUSTER_NAME,
            count=NUM_TASKS_TO_RUN,
            launchType='FARGATE',
            networkConfiguration={
                'awsvpcConfiguration': {
                    'subnets': subnet_ids,
                    'securityGroups': [
                        security_group_id,
                    ],
                    'assignPublicIp': 'DISABLED'
                }
            },
            overrides={
                'containerOverrides': container_overrides,
            },
            startedBy=TASK_STARTED_BY,
            tags=[
                {
                    'key': 'environment',
                    'value': ENVIRONMENT
                },
                {
                    'key': 'service_name',
                    'value': SERVICE_NAME
                },
            ],
            # Always run the latest task definition
            taskDefinition=CLUSTER_NAME
        )
        # If starting more than one task, the first should be representative
        if response['tasks']:
            task_arn = response['tasks'][0]['taskArn']
            return task_arn
        elif response['failures']:
            logging.error('Failures returned when starting task: {}'.format(
                response['failures']))
            raise SystemExit('Error running task')
    except ClientError as error:
        logging.error(error.response['Error']['Message'])
        raise SystemExit('Error running task')


def main():
    """Main entrypoint function."""
    parser = argparse.ArgumentParser()
    parser.add_argument('-e', '--environment', required=False,
                        help='Comma-separated list of environment variables '
                        'to override in Name=Value pairs, '
                        'e.g. ENV=qa,CACHE_STATUS=disabled')
    parser.add_argument(
        '-c', '--command', required=False,
        help='Comma-separated list of command arguments')
    environment_variable_list = []
    commands_list = []
    args = parser.parse_args()
    if args.environment:
        for variable in [variable.split(
                         '=') for variable in args.environment.split(',')]:
            environment_variable_list.append(
                {'name': variable[0], 'value': variable[1]})
    if args.command:
        for item in args.command.split(','):
            commands_list.append(item)

    ec2_client = boto3.client('ec2', region_name=AWS_REGION)
    ecs_client = boto3.client('ecs', region_name=AWS_REGION)

    vpc_id = get_vpc_id(ec2_client)
    security_group_id = get_security_group(ec2_client, vpc_id)
    subnet_ids = get_subnets(ec2_client, vpc_id)
    logging.info('Subnet IDs are {} and security group is {}'.format(
        subnet_ids, security_group_id))
    task_arn = run_task(ecs_client,
        security_group_id, subnet_ids, environment_variable_list, commands_list)
    if not SKIP_HEALTHCHECK:
        poll_result = _verify_task(
            client=ecs_client,
            verify_mode=verify_mode,
            cluster_name=CLUSTER_NAME,
            task_arn=task_arn,
            container_name=CONTAINER_NAME,
            timeout=time.time() + TASK_POLLING_TIMEOUT,
            grace_period=TASK_GRACE_PERIOD
        )
        if poll_result is False:
            raise SystemExit(
                'Deployment failed. Check task logs for {}'.format(
                    SERVICE_NAME))
        else:
            logging.info(
                'Started running task {} successfully'.format(task_arn))
    else:
        logging.info('Task started, but healthcheck is being skipped.')


if __name__ == "__main__":
    main()