import argparse
import csv
import datetime
import os
import tempfile
import json

import boto3


def main() -> None:
    client = boto3.client("ecs")
    args = get_args()
    print(args)

    # find all cluster names based on cli args
    env = "prod"
    loc_id = "13"
    cluster_names = []
    for cluster_type in args.types:
        for content_type in args.content:
            for priority in args.priorities:
                cluster_names.append(
                    f"{env}-vector-{cluster_type}-{content_type}-phys-loc-{loc_id}-p{priority}"
                )

    # gather details about all tasks
    tasks = []
    for cluster_name in cluster_names:
        print(cluster_name)
        tasks += get_tasks(client, cluster_name)
    print(f"Found {len(tasks)} tasks in total")

    # write results to file
    with open(
        f"{tempfile.gettempdir()}{os.path.sep}{args.filename}.{args.format}", "w"
    ) as f:
        if args.format == "json":
            json.dump(tasks, f, indent=4, default=str)
        elif args.format == "csv":
            fields = [
                "clusterArn",
                "age",
                "createdAt",
                "lastStatus",
                "healthStatus",
                "connectivity",
                "connectivityAt",
                "stopCode",
                "stoppedAt",
                "stoppedReason",
                "taskArn",
            ]
            writer = csv.DictWriter(f, fieldnames=fields)
            writer.writeheader()
            for task in tasks:
                writer.writerow({field: task.get(field, "") for field in fields})


def get_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )
    parser.add_argument("filename", type=str, help="output filename without extention")
    parser.add_argument(
        "--format",
        type=str,
        choices=["json", "csv"],
        default="json",
        help="output file format",
    )
    priority_choices = list(range(1, 7))
    parser.add_argument(
        "--priorities",
        type=int,
        nargs="+",
        choices=priority_choices,
        default=priority_choices,
        help="list of priorities to filter tasks by",
    )
    content_type_choices = ["audio", "video", "physical"]
    parser.add_argument(
        "--content",
        type=str,
        nargs="+",
        choices=content_type_choices,
        default=content_type_choices,
        help="list of task content types to filter tasks by",
    )
    type_choices = ["delivery", "encoder"]
    parser.add_argument(
        "--types",
        type=str,
        nargs="+",
        choices=type_choices,
        default=type_choices,
        help="list of task types to filter tasks by",
    )
    return parser.parse_args()


def get_tasks(client: boto3.client, cluster_name: str) -> list[dict]:
    return get_task_details(client, cluster_name, get_task_arns(client, cluster_name))


def get_task_arns(client: boto3.client, cluster_name: str) -> list[str]:
    paginator = client.get_paginator("list_tasks")
    page_iterator = paginator.paginate(cluster=cluster_name)
    task_arns = []
    for page in page_iterator:
        task_arns += page["taskArns"]
    return task_arns


def get_task_details(
    client: boto3.client, cluster_name: str, task_arns: list
) -> list[dict]:
    chunk_size = 100
    chunks = [
        task_arns[i : i + chunk_size] for i in range(0, len(task_arns), chunk_size)
    ]
    tasks = []
    for chunk in chunks:
        response = client.describe_tasks(cluster=cluster_name, tasks=chunk)
        tasks += response["tasks"]
    for task in tasks:
        task["age"] = datetime.datetime.now().astimezone() - task["createdAt"]
    return tasks


if __name__ == "__main__":
    main()