import boto3, sys


def _get_boto_session(region_name: str) -> boto3.Session:
    """
    Get the boto3 session. Works with profile configured in runtime.

    :param region_name: The configured AWS profile name.
    :returns: The Boto session.
    """
    if region_name == None:
        session = boto3.Session(region_name="us-west-2")
        return session
    else:
        session = boto3.Session(region_name=region_name)
        return session


def _get_node_instance_ids(
    session: boto3.Session,
    trainer_name: str,
    student_count: int,
    node_role: str,
    nodes_per_student: int,
) -> list:
    node_count_target = nodes_per_student * student_count
    ec2_client = session.client("ec2")
    instance_ids = []
    response = ec2_client.describe_instances(
        Filters=[
            {
                "Name": "tag:Name",
                "Values": [
                    node_role,
                ],
            },
            {
                "Name": "instance-state-name",
                "Values": [
                    "running",
                ],
            },
            {
                "Name": "tag:trainer",
                "Values": [
                    trainer_name,
                ],
            },
        ],
        DryRun=False,
        MaxResults=500,
    )
    node_count_actual = 0
    for reservation in response["Reservations"]:
        for instances in reservation["Instances"]:
            node_count_actual += 1
            instance_ids += [instances["InstanceId"]]
    while "NextToken" in response and response["NextToken"] != None:
        print("Fetching EC2 instances information with the next token.")
        nt = response["NextToken"]
        response = ec2_client.describe_instances(
            Filters=[
                {
                    "Name": "tag:Name",
                    "Values": [
                        node_role,
                    ],
                },
                {
                    "Name": "instance-state-name",
                    "Values": [
                        "running",
                    ],
                },
                {
                    "Name": "tag:trainer",
                    "Values": [
                        trainer_name,
                    ],
                },
            ],
            DryRun=False,
            MaxResults=500,
            NextToken=nt,
        )
        for reservation in response["Reservations"]:
            for instances in reservation["Instances"]:
                node_count_actual += 1
                instance_ids += [instances["InstanceId"]]
    print(
        str(student_count)
        + " students with "
        + str(node_count_actual)
        + " "
        + node_role
        + " node(s)."
    )
    if node_count_actual != node_count_target:
        print(
            "ERROR: Total available "
            + node_role
            + " nodes count is "
            + str(node_count_actual)
            + " instead of "
            + str(node_count_target)
            + "."
        )
        print("SOLUTION: Please wait and retry later or contact PE team.")
        exit(1)
    return instance_ids


def _assign_nodes_to_students(
    session: boto3.Session,
    instance_ids: list,
    node_role: str,
    nodes_per_student: int,
):
    """
    Nodes are tagged with `Name:{node_role}` by ASG and Launch Template.
    This routine will find all {node_role} nodes and then add new tag `role` to
        `role:{node_role}1`,
        `role:{node_role}2`,
        `role:{node_role}3` respectively based on student count.
    And, it depends on `nodes_per_student`.
    """
    ec2_client = session.client("ec2")
    student_idx = 1
    student_node_idx = 1
    for node_instance_id in instance_ids:
        print(
            "Tagging "
            + node_instance_id
            + " with class-user"
            + str(student_idx)
            + ":"
            + node_role
            + str(student_node_idx)
        )
        ec2_client.create_tags(
            DryRun=False,
            Resources=[
                node_instance_id,
            ],
            Tags=[
                {
                    "Key": "role",
                    "Value": node_role + str(student_node_idx),
                },
                {
                    "Key": "student",
                    "Value": "user" + str(student_idx),
                },
            ],
        )
        student_node_idx += 1
        if student_node_idx > nodes_per_student:
            student_node_idx = 1
            student_idx += 1
    print("Tagging " + node_role + " nodes completed.")


def _get_monitor_instance_ids(
    session: boto3.Session, trainer_name: str, student_count: int, region_name: str
) -> list:
    return _get_node_instance_ids(
        session, trainer_name, student_count, region_name + "-monitor", 1
    )


def _get_app_instance_ids(
    session: boto3.Session, trainer_name: str, student_count: int, region_name: str
) -> list:
    return _get_node_instance_ids(
        session, trainer_name, student_count, region_name + "-app", 1
    )


def _get_pd_instance_ids(
    session: boto3.Session, trainer_name: str, student_count: int, region_name: str
) -> list:
    return _get_node_instance_ids(
        session, trainer_name, student_count, region_name + "-pd", 1
    )


def _get_kv_instance_ids(
    session: boto3.Session, trainer_name: str, student_count: int, region_name: str
) -> list:
    return _get_node_instance_ids(
        session, trainer_name, student_count, region_name + "-kv", 1
    )


def _get_db_instance_ids(
    session: boto3.Session, trainer_name: str, student_count: int, region_name: str
) -> list:
    return _get_node_instance_ids(
        session, trainer_name, student_count, region_name + "-db", 1
    )


def _assign_pd_to_students(
    session: boto3.Session, instance_ids: list, region_f_name: str
):
    _assign_nodes_to_students(session, instance_ids, region_f_name + "-pd", 1)


def _assign_kv_to_students(
    session: boto3.Session, instance_ids: list, region_f_name: str
):
    _assign_nodes_to_students(session, instance_ids, region_f_name + "-kv", 1)


def _assign_db_to_students(
    session: boto3.Session, instance_ids: list, region_f_name: str
):
    _assign_nodes_to_students(session, instance_ids, region_f_name + "-db", 1)


def _assign_monitor_to_students(
    session: boto3.Session, instance_ids: list, region_f_name: str
):
    _assign_nodes_to_students(session, instance_ids, region_f_name + "-monitor", 1)


def _assign_app_to_students(
    session: boto3.Session, instance_ids: list, region_f_name: str
):
    _assign_nodes_to_students(session, instance_ids, region_f_name + "-app", 1)


if __name__ == "__main__":
    student_count = sys.argv[1]
    trainer_name = sys.argv[2]
    region_name = sys.argv[3]
    region_f_name = sys.argv[4]
    session = _get_boto_session(region_name)

    if region_name == "us-west-2":
        monitor_instance_ids = _get_monitor_instance_ids(
            session, trainer_name, int(student_count), region_name
        )
        pd_instance_ids = _get_pd_instance_ids(
            session, trainer_name, int(student_count), region_name
        )

        kv_instance_ids = _get_kv_instance_ids(
            session, trainer_name, int(student_count), region_name
        )
        _assign_monitor_to_students(session, monitor_instance_ids, region_f_name)
        _assign_pd_to_students(session, pd_instance_ids, region_f_name)
        _assign_kv_to_students(session, kv_instance_ids, region_f_name)
    else:
        app_instance_ids = _get_app_instance_ids(
            session, trainer_name, int(student_count), region_name
        )
        db_instance_ids = _get_db_instance_ids(
            session, trainer_name, int(student_count), region_name
        )
        pd_instance_ids = _get_pd_instance_ids(
            session, trainer_name, int(student_count), region_name
        )

        kv_instance_ids = _get_kv_instance_ids(
            session, trainer_name, int(student_count), region_name
        )
        _assign_app_to_students(session, app_instance_ids, region_f_name)
        _assign_pd_to_students(session, pd_instance_ids, region_f_name)
        _assign_kv_to_students(session, kv_instance_ids, region_f_name)
        _assign_db_to_students(session, db_instance_ids, region_f_name)