import argparse
import json
import time

import boto3

ecs_client = boto3.client('ecs', region_name='us-east-1')
iam_client = boto3.client('iam')

ecs_exec_policy = {
    "Version": "2012-10-17",
    "Statement": [
        {
            "Effect": "Allow",
            "Action": [
                "ssmmessages:CreateControlChannel",
                "ssmmessages:CreateDataChannel",
                "ssmmessages:OpenControlChannel",
                "ssmmessages:OpenDataChannel"
            ],
            "Resource": "*"
        }
    ]
}


def main():
    """Enable ECS exec for a service."""
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '-a', '--action', required=False, default='enable',
        help='Action to take: enable or disable'
    )
    parser.add_argument(
        '-s', '--service', required=True,
        help='Name of service, including environment '
        'e.g. qa-ows-track, prod-graphql-product'
    )
    args = parser.parse_args()
    
    paginator = iam_client.get_paginator('list_policies')
    response_iterator = paginator.paginate(
        Scope='Local',
        OnlyAttached=False,
        PathPrefix='/',
        PolicyUsageFilter='PermissionsPolicy',
        PaginationConfig={'PageSize': 100},
    )

    for response in response_iterator:
        policy_arn = next((policy['Arn'] for policy in response['Policies']
                        if policy['PolicyName'] == 'temp-ecs-exec-policy'),
                        None)
        if policy_arn:
            print('A policy named temp-ecs-exec-policy already exists')
            break
    if args.action == 'enable':
        if not policy_arn:
            print('Creating new policy named temp-ecs-exec-policy')
            policy = iam_client.create_policy(
                PolicyName='temp-ecs-exec-policy',
                PolicyDocument=json.dumps(ecs_exec_policy),
                Description='temp-ecs-exec-policy',
                Tags=[
                    {
                        'Key': 'environment',
                        'Value': 'prod'
                    },
                    {
                        'Key': 'service_name',
                        'Value': 'ecs-exec'
                    },
                    {
                        'Key': 'temporary',
                        'Value': 'true'
                    },
                ]
            )
            policy_arn = policy['Policy']['Arn']
        print(f'Attaching temp-ecs-exec-policy to {args.service}-task-role')
        iam_client.attach_role_policy(
            RoleName=f'{args.service}-task-role',
            PolicyArn=policy_arn,
        )
    else:
        print(f'Detaching temp-ecs-exec-policy from {args.service}-task-role')
        iam_client.detach_role_policy(
            RoleName=f'{args.service}-task-role',
            PolicyArn=policy_arn,
        )
        
        print(f'Deleting {policy_arn}')
        iam_client.delete_policy(
            PolicyArn=policy_arn
        )

    latest_task = ecs_client.describe_task_definition(
        taskDefinition=args.service)
    # Just get the task definition element
    task_definition_arn = latest_task['taskDefinition']['taskDefinitionArn']
    
    if args.action == 'enable':
        enable_ecs_exec = True
    else:
        enable_ecs_exec = False

    update = ecs_client.update_service(
        cluster=args.service,
        service=args.service,
        taskDefinition=task_definition_arn,
        enableExecuteCommand=enable_ecs_exec,
        forceNewDeployment=True,
    )

    print(f"Service response: {update}")
    print(f"Service status: {update['service']['status']}")

    # Let's be lazy
    print('Sleeping for 5 minutes to allow deployment to complete.')
    time.sleep(300)

    container_name = '-'.join(args.service.split('-')[1:])
    task_response = ecs_client.list_tasks(
        cluster=args.service, serviceName=args.service)
    for index, task in enumerate(task_response['taskArns']):
        print(f'New task ARN number {index}: {task}')
        if args.action == 'enable':
            print(f'Now run this: aws ecs execute-command --region us-east-1 --cluster {args.service} --task {task} --container {container_name} --command "sh" --interactive') # noqa


if __name__ == "__main__":
    main()