#!/usr/bin/env python3

import boto3
import logging
import os

from typing import List, Tuple


# Pricing details per hour
CPU_PRICE_PER_HOUR=0.04048
MEMORY_PRICE_PER_HOUR=0.004445

ARM_CPU_PRICE_PER_HOUR=0.03238
ARM_MEMORY_PRICE_PER_HOUR=0.00356

# Hours in a month (approx)
HOURS_PER_MONTH = 730

def calculate_fargate_cost(client, cluster):
    response = client.list_tasks(
        cluster=cluster,
        desiredStatus='RUNNING'
    )

    x86_total_cost_per_month = arm_total_cost_per_month = 0.0

    if len(response['taskArns']) > 0:
        task_descriptions = client.describe_tasks(
            cluster=cluster,
            tasks=response['taskArns']
        )

        total_cpu = 0.0
        total_memory = 0.0

        for task in task_descriptions['tasks']:
            logging.debug(f"{task['cpu']=} CPU")
            total_cpu += float(task['cpu']) / 1024
            logging.debug(f"{total_cpu} CPU")
            logging.debug(f"{task['memory']=} GB")
            total_memory += float(task['memory']) / 1024
            logging.debug(f"{total_memory=} GB")

        # AWS charges per vCPU per hour and per GB of memory per hour, so:
        x86_cpu_cost_per_month = total_cpu * CPU_PRICE_PER_HOUR * HOURS_PER_MONTH
        arm_cpu_cost_per_month = total_cpu * ARM_CPU_PRICE_PER_HOUR * HOURS_PER_MONTH
        logging.debug(f"x86 {x86_cpu_cost_per_month=} and arm {arm_cpu_cost_per_month=}")
        x86_memory_cost_per_month = total_memory * MEMORY_PRICE_PER_HOUR * HOURS_PER_MONTH
        arm_memory_cost_per_month = total_memory * ARM_MEMORY_PRICE_PER_HOUR * HOURS_PER_MONTH
        logging.debug(f"x86 {x86_memory_cost_per_month=} and arm {arm_memory_cost_per_month=}")

        x86_total_cost_per_month = x86_cpu_cost_per_month + x86_memory_cost_per_month
        arm_total_cost_per_month = arm_cpu_cost_per_month + arm_memory_cost_per_month

    return x86_total_cost_per_month, arm_total_cost_per_month


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_profiles() -> List:
    return [
        'gdb-apollo-dev',
        'gdb-apollo-prod',
        # 'gdb-artistapp-dev',
        # 'gdb-artistapp-prod',
        # 'gdb-databricks-dev',
        # 'gdb-databricks-prod',
        # 'gdb-delphi-dev',
        # 'gdb-delphi-prod',
        # 'gdb-infra-dev',
        # 'gdb-infra-prod',
        # 'gdb-mct-dev',
        # 'gdb-mct-prod',
        # 'gdb-datasci-dev',
        # 'gdb-core-dev',
        # 'gdb-core-prod',
        # 'gdb-whitelist-legacy',
        # 'gdb-whitelist-dev',
        # 'gdb-whitelist-prod',
        # 'gdb-smu-dev',
        # 'gdb-smu-prod'
    ]


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

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

    for key in ecs_clients:
        clusters = get_clusters(ecs_clients[key])
        for cluster in clusters:
            x86_cost, arm_cost = calculate_fargate_cost(ecs_clients[key], cluster)
            logging.info(f"The estimated cost for running tasks on ECS Fargate in {cluster=} per month: {x86_cost} USD on x86 vs {arm_cost} USD on ARM")


if __name__ == '__main__':
    logging.basicConfig(level=os.environ.get("LOGLEVEL", "INFO"))
    main()