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 _run_nodes(
    session: boto3.Session,
    trainer_name: str,
    student_count: int,
    node_type: str,
    nodes_per_student: int,
):
    nodes_count_target = nodes_per_student * student_count
    asg_client = session.client("autoscaling")
    asg_name = "classroom-" + trainer_name + "-" + node_type + "-asg"
    asg_client.update_auto_scaling_group(
        AutoScalingGroupName=asg_name,
        MinSize=nodes_count_target,
        MaxSize=nodes_count_target,
        DesiredCapacity=nodes_count_target,
    )


def _run_pd_nodes(session: boto3.Session, trainer_name: str, student_count: int):
    _run_nodes(session, trainer_name, student_count, "pd", 3)


def _run_kv_nodes(session: boto3.Session, trainer_name: str, student_count: int):
    _run_nodes(session, trainer_name, student_count, "kv", 3)


def _run_db_nodes(session: boto3.Session, trainer_name: str, student_count: int):
    _run_nodes(session, trainer_name, student_count, "db", 1)


def _run_tiflash_nodes(session: boto3.Session, trainer_name: str, student_count: int):
    _run_nodes(session, trainer_name, student_count, "tiflash", 1)


if __name__ == "__main__":

    student_count = int(sys.argv[1])
    if student_count > 14:
        exit(1)
    trainer_name = sys.argv[2]
    session = _get_boto_session(None)
    _run_pd_nodes(session, trainer_name, student_count)
    _run_kv_nodes(session, trainer_name, student_count)
    _run_db_nodes(session, trainer_name, student_count)
    _run_tiflash_nodes(session, trainer_name, student_count)