import boto3, sys


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

    :param profile_name: The configured AWS profile name.
    :returns: The Boto session.
    """
    if profile_name == None:
        session = boto3.Session(region_name="us-west-2")
        return session
    else:
        session = boto3.Session(profile_name=profile_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_az1monitor_instance_ids(
    session: boto3.Session, trainer_name: str, student_count: int
) -> list:
    return _get_node_instance_ids(session, trainer_name, student_count, "az1monitor", 1)


def _get_az2monitor_instance_ids(
    session: boto3.Session, trainer_name: str, student_count: int
) -> list:
    return _get_node_instance_ids(session, trainer_name, student_count, "az2monitor", 1)


def _get_az1pd_instance_ids(
    session: boto3.Session, trainer_name: str, student_count: int
) -> list:
    return _get_node_instance_ids(session, trainer_name, student_count, "az1pd", 2)


def _get_az2pd_instance_ids(
    session: boto3.Session, trainer_name: str, student_count: int
) -> list:
    return _get_node_instance_ids(session, trainer_name, student_count, "az2pd", 2)


def _get_az1kv_instance_ids(
    session: boto3.Session, trainer_name: str, student_count: int
) -> list:
    return _get_node_instance_ids(session, trainer_name, student_count, "az1kv", 2)


def _get_az2kv_instance_ids(
    session: boto3.Session, trainer_name: str, student_count: int
) -> list:
    return _get_node_instance_ids(session, trainer_name, student_count, "az2kv", 2)


def _get_az1db_instance_ids(
    session: boto3.Session, trainer_name: str, student_count: int
) -> list:
    return _get_node_instance_ids(session, trainer_name, student_count, "az1db", 1)


def _get_az2db_instance_ids(
    session: boto3.Session, trainer_name: str, student_count: int
) -> list:
    return _get_node_instance_ids(session, trainer_name, student_count, "az2db", 1)


def _assign_az1pd_to_students(session: boto3.Session, instance_ids: list):
    _assign_nodes_to_students(session, instance_ids, "az1pd", 2)


def _assign_az2pd_to_students(session: boto3.Session, instance_ids: list):
    _assign_nodes_to_students(session, instance_ids, "az2pd", 2)


def _assign_az1kv_to_students(session: boto3.Session, instance_ids: list):
    _assign_nodes_to_students(session, instance_ids, "az1kv", 2)


def _assign_az2kv_to_students(session: boto3.Session, instance_ids: list):
    _assign_nodes_to_students(session, instance_ids, "az2kv", 2)


def _assign_az1db_to_students(session: boto3.Session, instance_ids: list):
    _assign_nodes_to_students(session, instance_ids, "az1db", 1)


def _assign_az2db_to_students(session: boto3.Session, instance_ids: list):
    _assign_nodes_to_students(session, instance_ids, "az2db", 1)


def _assign_az1monitor_to_students(session: boto3.Session, instance_ids: list):
    _assign_nodes_to_students(session, instance_ids, "az1monitor", 1)


def _assign_az2monitor_to_students(session: boto3.Session, instance_ids: list):
    _assign_nodes_to_students(session, instance_ids, "az2monitor", 1)


if __name__ == "__main__":
    student_count = sys.argv[1]
    trainer_name = sys.argv[2]
    session = _get_boto_session(None)
    az1monitor_instance_ids = _get_az1monitor_instance_ids(
        session, trainer_name, int(student_count)
    )
    az2monitor_instance_ids = _get_az2monitor_instance_ids(
        session, trainer_name, int(student_count)
    )
    az1pd_instance_ids = _get_az1pd_instance_ids(
        session, trainer_name, int(student_count)
    )
    az2pd_instance_ids = _get_az2pd_instance_ids(
        session, trainer_name, int(student_count)
    )
    az1kv_instance_ids = _get_az1kv_instance_ids(
        session, trainer_name, int(student_count)
    )
    az2kv_instance_ids = _get_az2kv_instance_ids(
        session, trainer_name, int(student_count)
    )
    az1db_instance_ids = _get_az1db_instance_ids(
        session, trainer_name, int(student_count)
    )
    az2db_instance_ids = _get_az2db_instance_ids(
        session, trainer_name, int(student_count)
    )

    _assign_az1monitor_to_students(session, az1monitor_instance_ids)
    _assign_az2monitor_to_students(session, az2monitor_instance_ids)
    _assign_az1pd_to_students(session, az1pd_instance_ids)
    _assign_az2pd_to_students(session, az2pd_instance_ids)
    _assign_az1kv_to_students(session, az1kv_instance_ids)
    _assign_az2kv_to_students(session, az2kv_instance_ids)
    _assign_az1db_to_students(session, az1db_instance_ids)
    _assign_az2db_to_students(session, az2db_instance_ids)