#!/usr/bin/env python3

from awsaccountmgr import (
    read_config_files,
    Organization,
    delete_default_vpc,
    create_boto3_client,
    assume_role,
    Account,
)
import logging
import boto3
from argparse import ArgumentParser
from concurrent.futures import ThreadPoolExecutor


def main():
    # Parse/retrieve required parameters
    root_ou, config_folder, logging_level = parse_args()
    master_account_id = Organization.get_master_account_id()

    # Set logging
    logging_format = "%(asctime)-15s %(levelname)s %(message)s"
    logging.basicConfig(format=logging_format, level=logging_level)

    # Initialize Organization class
    if boto3.client("sts").get_caller_identity().get("Account") != master_account_id:
        logging.info(
            f"Script not running from master account, assuming OrganizationAccountAccessRole"
            f" in master account {master_account_id}"
        )
        credentials = assume_role(master_account_id)
    else:
        session = boto3.Session()
        current_creds = session.get_credentials().get_frozen_credentials()
        credentials = (
            current_creds.access_key,
            current_creds.secret_key,
            current_creds.token,
        )

    organization = Organization(root_ou, credentials)

    # Parse config folder
    accounts = read_config_files(config_folder)
    logging.info(f"Found {len(accounts)} account(s) in configuration file.")

    # Create/Update accounts
    for account in accounts:
        create_or_update_account(organization, account)


def create_or_update_account(org_session, account: Account):
    """Creates or updates a single AWS account.

    :param org_session: Instance of Organization class
    :param account: Instance of Account class
    """
    # Create new account
    account_id = org_session.get_account_id(account.full_name)

    if not account_id:
        logging.info(f"Creating new account {account.full_name}")
        account_id = org_session.create_account(account)

    # Move account to given OU
    logging.info(f"Moving account {account_id} to OU {account.ou_path}")
    org_session.move_account(
        account_id, account.ou_path, account.allow_direct_move_between_ou
    )

    # Remove default VPC
    if account.delete_default_vpc is True:

        # Retrieve all regions
        ec2_client = boto3.client("ec2")
        all_regions = [
            region["RegionName"]
            for region in ec2_client.describe_regions(AllRegions=False)["Regions"]
        ]

        # Remove VPCs from all regions using threads
        args = ((account_id, org_session, region) for region in all_regions)
        with ThreadPoolExecutor(max_workers=10) as executor:
            for _ in executor.map(lambda f: schedule_delete_default_vpc(*f), args):
                pass

    # Create account alias
    logging.info(f"Creating/updating account {account_id} alias {account.alias}")
    org_session.create_account_alias(account_id, account.alias)

    # Create/Update Alternate Contacts
    if account.update_alternate_contacts:
        if account.operations_contact:
            logging.info(f"Updating operations contact for account {account_id}")
            org_session.update_alternate_contact(
                account_id, "OPERATIONS", account.operations_contact
            )
        else:
            logging.info(f"Removing operations contact from {account_id}")
            org_session.remove_alternate_contact(account_id, "OPERATIONS")

        if account.security_contact:
            logging.info(f"Updating security contact for account {account_id}")
            org_session.update_alternate_contact(
                account_id, "SECURITY", account.security_contact
            )
        else:
            logging.info(f"Removing security contact from {account_id}")
            org_session.remove_alternate_contact(account_id, "SECURITY")

        if account.billing_contact:
            logging.info(f"Updating billing contact for account {account_id}")
            org_session.update_alternate_contact(
                account_id, "BILLING", account.billing_contact
            )
        else:
            logging.info(f"Removing billing contact from {account_id}")
            org_session.remove_alternate_contact(account_id, "BILLING")

    # Add tags
    if account.tags:
        logging.info(f"Adding tags to account {account_id}: {account.tags}")
        org_session.create_account_tags(account_id, account.tags)


def schedule_delete_default_vpc(account_id, org_session, region):
    """Schedule a delete_default_vpc on a thread

    :param account_id: The account ID to remove the VPC from
    :param org_session: The Organization class instance
    :param region: The name of the region the VPC is resided
    """
    # Remove VPC in given region
    ec2_client = create_boto3_client(
        account_id, "ec2", org_session._master_credentials, region_name=region
    )
    logging.info(f"Deleting default VPC from {account_id} in region {region}")
    delete_default_vpc(ec2_client, account_id)


def parse_args():
    """Parse command line arguments"""
    parser = ArgumentParser(
        description="Multi account creation and management in AWS Organization."
    )

    parser.add_argument(
        "root_ou", help="The ID of the root Organizational Unit (eg. r-abc1)"
    )

    parser.add_argument(
        "config_folder", help="The folder containing the account configuration files"
    )

    parser.add_argument(
        "--logging-level",
        default="INFO",
        help="The logging output level, defaults to INFO",
    )

    args = parser.parse_args()
    return (args.root_ou, args.config_folder, args.logging_level)


if __name__ == "__main__":
    main()
