#!/usr/bin/env python3

import yaml
from awsaccountmgr import read_config_files, Organization, delete_default_vpc, create_boto3_client, assume_role
import logging
import os
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):
    """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)

    # 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()
