import csv
from typing import List

from awsume.awsumepy import awsume
import boto3
from tqdm import tqdm

from packages.python.accounts.gdb import ACCOUNTS as gdb_accounts
from packages.python.accounts.theorchard import ACCOUNTS as theorchard_accounts


def get_session(account_alias, use_awsume=False):
    if use_awsume:
        session = awsume(account_alias)
    else:
        session = boto3.Session(profile_name=account_alias)

    return session


def get_clusters(client) -> List:
    """
    return a list of cluster arns
    """
    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:
    """
    return a list of service arns for a given cluster
    """
    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 a list of task arns for a given service
    """
    paginator = client.get_paginator('list_tasks')
    page_iterator = paginator.paginate(cluster=cluster, 
                                       serviceName=service, 
                                       launchType='FARGATE',
                                       desiredStatus='RUNNING'
                                       )
    return [item for sublist in page_iterator for item in sublist['taskArns']]


def get_desire_count(client, cluster: str, service: str) -> int:
    """
    return the desired count for a given service
    """
    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:
    """
    return the max count for a given service

    This is the max count of task for a given service with autoscaling enabled
    """
    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 extract_cluster_name_from_arn(arn: str) -> str:
    """
    return the cluster name from the arn
    """
    return arn.split('/')[-1]


def extract_service_name_from_arn(arn: str) -> str:
    """
    return the service name from the arn
    """
    return arn.split('/')[-1]


def main():
    """
    Get the number of clusters, services, tasks, and max number of tasks for each account
    """
    gdb_ecs_clients = {}
    gdb_autoscaling_clients = {}

    theorchard_ecs_clients = {}
    theorchard_autoscaling_clients = {}

    statistics = {}
    statistics_per_service = {}
    statistics_per_family = {}

    total_number_of_services = {}
    total_number_of_tasks = {}
    max_number_of_tasks = {}

    for account_alias in theorchard_accounts.values():
        total_number_of_services[account_alias] = 0
        total_number_of_tasks[account_alias] = 0
        max_number_of_tasks[account_alias] = 0

        # Create boto3 session
        session = get_session(account_alias, True)
        # Create ecs client
        theorchard_ecs_clients[account_alias] = session.client('ecs')
        # Create autoscaling client
        theorchard_autoscaling_clients[account_alias] = session.client('application-autoscaling')

    for account_alias, client in theorchard_ecs_clients.items():
        print(f'getting items for {account_alias}')
        # Get list of clusters
        clusters = get_clusters(client)
        # Get list of services
        for cluster in clusters:
            print(f'getting items for {cluster}')
            services = get_services(client, cluster)
            total_number_of_services[account_alias] += len(services)
            # Get list of tasks
            for service in tqdm(services, desc="Processing services"):
                tasks = get_list_of_tasks(client, cluster, service)
                tasks_count = len(tasks)
                total_number_of_tasks[account_alias] += tasks_count
                autoscaling_client = theorchard_autoscaling_clients[account_alias]
                max_count = get_max_count(autoscaling_client, cluster, service)
                max_number_of_tasks[account_alias] += max_count if max_count > 0 else tasks_count

                service_name = extract_service_name_from_arn(service)
                environment, family, *name = service_name.split('-')
                service_name = '-'.join(name)

                if family not in statistics_per_family:
                    statistics_per_family[family] = {'number_of_services': 0, 'number_of_tasks': 0, 'max_number_of_tasks': 0}

                if (environment, family) not in statistics:
                    statistics[(environment, family)] = {'number_of_services': 0, 'number_of_tasks': 0, 'max_number_of_tasks': 0}

                if (environment, family, service_name) not in statistics_per_service:
                    statistics_per_service[(environment, family, service_name)] = {'number_of_services': 0, 'number_of_tasks': 0, 'max_number_of_tasks': 0}

                statistics[(environment, family)]['number_of_services'] += 1
                statistics[(environment, family)]['number_of_tasks'] += tasks_count
                statistics[(environment, family)]['max_number_of_tasks'] += max_count if max_count > 0 else tasks_count

                statistics_per_service[(environment, family, service_name)]['number_of_services'] += 1
                statistics_per_service[(environment, family, service_name)]['number_of_tasks'] += tasks_count
                statistics_per_service[(environment, family, service_name)]['max_number_of_tasks'] += max_count if max_count > 0 else tasks_count

                statistics_per_family[family]['number_of_services'] += 1
                statistics_per_family[family]['number_of_tasks'] += tasks_count
                statistics_per_family[family]['max_number_of_tasks'] += max_count if max_count > 0 else tasks_count

    for account_alias in gdb_accounts.values():
        total_number_of_services[account_alias] = 0
        total_number_of_tasks[account_alias] = 0
        max_number_of_tasks[account_alias] = 0

        # Create boto3 session
        session = get_session(account_alias, False)
        # Create ecs client
        gdb_ecs_clients[account_alias] = session.client('ecs')
        # Create autoscaling client
        gdb_autoscaling_clients[account_alias] = session.client('application-autoscaling')

    for account_alias, client in gdb_ecs_clients.items():
        print(f'getting items for {account_alias}')
        # Get list of clusters
        clusters = get_clusters(client)
        # Get list of services
        for cluster in clusters:
            print(f'getting items for {cluster}')
            services = get_services(client, cluster)
            total_number_of_services[account_alias] += len(services)
            # Get list of tasks
            for service in tqdm(services, desc="Processing services"):
                tasks = get_list_of_tasks(client, cluster, service)
                tasks_count = len(tasks)
                total_number_of_tasks[account_alias] += tasks_count
                autoscaling_client = gdb_autoscaling_clients[account_alias]
                max_count = get_max_count(autoscaling_client, cluster, service)
                max_number_of_tasks[account_alias] += max_count if max_count > 0 else tasks_count

                service_name = extract_service_name_from_arn(service)
                environment, family, *name = service_name.split('-')
                service_name = '-'.join(name)

                if family not in statistics_per_family:
                    statistics_per_family[family] = {'number_of_services': 0, 'number_of_tasks': 0, 'max_number_of_tasks': 0}

                if (environment, family) not in statistics:
                    statistics[(environment, family)] = {'number_of_services': 0, 'number_of_tasks': 0, 'max_number_of_tasks': 0}

                if (environment, family, service_name) not in statistics_per_service:
                    statistics_per_service[(environment, family, service_name)] = {'number_of_services': 0, 'number_of_tasks': 0, 'max_number_of_tasks': 0}

                statistics[(environment, family)]['number_of_services'] += 1
                statistics[(environment, family)]['number_of_tasks'] += tasks_count
                statistics[(environment, family)]['max_number_of_tasks'] += max_count if max_count > 0 else tasks_count

                statistics_per_service[(environment, family, service_name)]['number_of_services'] += 1
                statistics_per_service[(environment, family, service_name)]['number_of_tasks'] += tasks_count
                statistics_per_service[(environment, family, service_name)]['max_number_of_tasks'] += max_count if max_count > 0 else tasks_count

                statistics_per_family[family]['number_of_services'] += 1
                statistics_per_family[family]['number_of_tasks'] += tasks_count
                statistics_per_family[family]['max_number_of_tasks'] += max_count if max_count > 0 else tasks_count
    

    # Write down totals per account to CSV
    with open('totals.csv', 'w', newline='') as csvfile:
        fieldnames = ['account', 'number_of_services', 'number_of_tasks', 'max_number_of_tasks']
        writer = csv.DictWriter(csvfile, fieldnames=fieldnames)

        # Write the headers to the CSV file
        writer.writeheader()

        # Write each row to the CSV file
        for account_alias in gdb_accounts.values():
            writer.writerow({**{'account': account_alias}, **{'number_of_services': total_number_of_services[account_alias], 'number_of_tasks': total_number_of_tasks[account_alias], 'max_number_of_tasks': max_number_of_tasks[account_alias]}})

        for account_alias in theorchard_accounts.values():
            writer.writerow({**{'account': account_alias}, **{'number_of_services': total_number_of_services[account_alias], 'number_of_tasks': total_number_of_tasks[account_alias], 'max_number_of_tasks': max_number_of_tasks[account_alias]}})
    
    # Open the CSV file
    with open('statistics.csv', 'w', newline='') as csvfile:
        fieldnames = ['environment', 'family', 'number_of_services', 'number_of_tasks', 'max_number_of_tasks']
        writer = csv.DictWriter(csvfile, fieldnames=fieldnames)

        # Write the headers to the CSV file
        writer.writeheader()

        # Write each row to the CSV file
        for (environment, family), stats in statistics.items():
            writer.writerow({**{'environment': environment, 'family': family}, **stats})

    # Open the CSV file
    with open('statistics-per-service.csv', 'w', newline='') as csvfile:
        fieldnames = ['environment', 'family', 'name', 'number_of_services', 'number_of_tasks', 'max_number_of_tasks']
        writer = csv.DictWriter(csvfile, fieldnames=fieldnames)

        # Write the headers to the CSV file
        writer.writeheader()

        # Write each row to the CSV file
        for (environment, family, service_name), stats in statistics_per_service.items():
            writer.writerow({**{'environment': environment, 'family': family, 'name': service_name}, **stats})

    # Open the CSV file
    with open('statistics-per-family.csv', 'w', newline='') as csvfile:
        fieldnames = ['family', 'number_of_services', 'number_of_tasks', 'max_number_of_tasks']
        writer = csv.DictWriter(csvfile, fieldnames=fieldnames)

        # Write the headers to the CSV file
        writer.writeheader()

        # Write each row to the CSV file
        for family, stats in statistics_per_family.items():
            writer.writerow({**{'family': family}, **stats})


if __name__ == '__main__':
    main()