#!/usr/bin/env python3

import boto3

from typing import List, Tuple


def get_clusters(client) -> List:
    paginator = client.get_paginator('list_clusters')
    page_iterator = paginator.paginate()
    return [item for sublist in page_iterator for item in
            sublist['clusterArns']]


def get_services(client, cluster: str) -> List:
    paginator = client.get_paginator('list_services')
    page_iterator = paginator.paginate(cluster=cluster, launchType='FARGATE')
    return [item for sublist in page_iterator for item in
            sublist['serviceArns']]


def get_list_of_tasks(client, cluster: str, service: str) -> List:
    """
    return list of task arns for a given service
    """
    paginator = client.get_paginator('list_tasks')
    page_iterator = paginator.paginate(cluster=cluster, serviceName=service,
                                       launchType='FARGATE')
    return [item for sublist in page_iterator for item in sublist['taskArns']]


def get_number_of_containers_for_task(client, cluster, task_arn):
    client = client.describe_tasks(
        cluster=cluster,
        tasks=[task_arn]
    )
    return len(client['tasks'][0]['containers'])


def get_desire_count(client, cluster: str, service: str) -> int:
    result = client.describe_services(cluster=cluster, services=[service])
    service = result['services'][0]
    return service['desiredCount']


def get_max_count(client, cluster: str, service: str) -> int:
    resource_id = f'service/{cluster.split("/")[-1]}/{service.split("/")[-1]}'
    result = client.describe_scalable_targets(
        ResourceIds=[resource_id],
        ScalableDimension='ecs:service:DesiredCount',
        ServiceNamespace='ecs',
    )
    if result['ScalableTargets']:
        scalable_targets = result['ScalableTargets'][0]
        return scalable_targets['MaxCapacity']

    return 0


def get_profiles() -> List:
    return [
        'gdb-apollo-dev',
        'gdb-apollo-prod',
        'gdb-artistapp-dev',
        'gdb-artistapp-prod',
        'gdb-core-dev',
        'gdb-core-prod',
        'gdb-databricks-dev',
        'gdb-databricks-prod',
        'gdb-datasci-dev',
        'gdb-delphi-dev',
        'gdb-delphi-prod',
        'gdb-infra-dev',
        'gdb-infra-prod',
        'gdb-mct-dev',
        'gdb-mct-prod',
        'gdb-prodhub-dev',
    ]


def main():
    profiles = get_profiles()
    sessions = []
    for profile in profiles:
        sessions.append(boto3.session.Session(profile_name=profile))

    ecs_clients = {}
    autoscaling_clients = {}
    for session in sessions:
        ecs_clients[session.profile_name] = session.client('ecs')
        autoscaling_clients[session.profile_name] = session.client(
            'application-autoscaling')

    total_count = 0
    max_capacity = 0
    total_number_of_containers = 0
    max_number_of_containers = 0
    for key in ecs_clients:
        clusters = get_clusters(ecs_clients[key])
        for cluster in clusters:
            services = get_services(ecs_clients[key], cluster)
            for service in services:
                tasks = get_list_of_tasks(
                    ecs_clients[key],
                    cluster,
                    service
                )
                containers_count = 0
                if len(tasks) > 0:
                    containers_count = get_number_of_containers_for_task(
                        ecs_clients[key],
                        cluster,
                        tasks[0]
                    )
                desired_count = get_desire_count(
                    ecs_clients[key],
                    cluster,
                    service
                )
                total_count += desired_count
                total_number_of_containers += desired_count * containers_count
                max_count = get_max_count(
                    autoscaling_clients[key],
                    cluster,
                    service
                )
                max_capacity += max_count
                max_number_of_containers += max_count * containers_count

    print(f'Total count of tasks across all accounts: {total_count}')
    print(f'Total count of containers across all accounts: {total_number_of_containers}')
    print(f'Max count of tasks across all accounts: {max_capacity}')
    print(f'Max count of containers across all accounts: {max_number_of_containers}')


if __name__ == '__main__':
    main()