#from botocore.exceptions import ClientError
import boto3
from pprint import pprint, pformat
#import json
#import sys
import argparse
import logging
import time

# Period in seconds to analyze data from log insights
ASSESSMENT_PERIOD = 3600



def parse_arguments():
    parser = argparse.ArgumentParser()
    parser.add_argument('--accesskey', help='AWS Access key', default='')
    parser.add_argument('--secretkey', help='AWS Secret key', default='')
    parser.add_argument('--profile', help='AWS profile', default='gdb-delphi-prod')
    parser.add_argument('--region', help='AWS Default region', default='us-east-1')
    parser.add_argument('--cluster', help='ECS cluster', default='prod-delphi-clu')
    parser.add_argument('--service', help='ECS service name', default='prod-delphi-api')
    return parser.parse_args()

def memory_query(taskId: str) -> str:
    return f"""fields @timestamp, @message
| filter (Type="Task" and TaskId="{taskId}")
| display MemoryUtilized/MemoryReserved*100
| sort @timestamp asc
| limit 1"""


if __name__ == "__main__":
    logger = logging.getLogger('')
    logger.setLevel(logging.DEBUG)

    args = parse_arguments()

    # Initialize the client with credentials or with IAM role.
    if (args.accesskey and args.secretkey):
        ecs_client = boto3.client(
          'ecs',
          aws_access_key_id = args.accesskey,
          aws_secret_access_key = args.secretkey,
          region_name = args.region
        )
        logs_client = boto3.client(
          'logs',
          aws_access_key_id = args.accesskey,
          aws_secret_access_key = args.secretkey,
          region_name = args.region
        )
    elif args.profile:
        session = boto3.Session(profile_name=args.profile,region_name = args.region)
        ecs_client = session.client('ecs')
        logs_client = session.client('logs')
    else:
        ecs_client = boto3.client('ecs', region_name=args.region)
        logs_client = boto3.client('logs', region_name=args.region)

    response = ecs_client.list_tasks(
        cluster=args.cluster,
        serviceName=args.service,
        desiredStatus='RUNNING'
    )

    tasks_ids = [arn.split('/')[-1] for arn in response['taskArns']]

    for task in tasks_ids:
        epoch_time = int(time.time())
        query_id = logs_client.start_query(
            logGroupName=f"/aws/ecs/containerinsights/{args.cluster}/performance",
            startTime = epoch_time - ASSESSMENT_PERIOD ,
            endTime = epoch_time,
            queryString=memory_query(task),
        )['queryId']
        response = logs_client.get_query_results(
            queryId=query_id
        )
        query_retry_limit = 100
        while response['status'] != 'Complete' or query_retry_limit <= 0:
            time.sleep(1)
            response = logs_client.get_query_results(
                queryId=query_id
            )
            query_retry_limit -= 1
        if query_retry_limit <= 0:
            print(f"TaskId: {task}, Logs Insights query timed out")
        print(f"TaskId: {task}, MemoryConsumption: {float(response['results'][0][0][('value')]):.1f}%")