import json
import argparse

import boto3

iam_client = boto3.client('iam')
ssm_client = boto3.client('ssm')


def generate_aws_cli_config(aws_accounts, source_profile_name):
    """Generate AWS CLI config for all accounts."""
    current_user = iam_client.get_user()
    username = current_user['User']['UserName']

    print("# AWS CLI Configuration for all accounts")
    print("# Add this to your ~/.aws/config file\n")

    for account in aws_accounts:
        print(f"[profile {account['name']}]")
        print(f"source_profile = {source_profile_name}")
        print(f"role_session_name = {username}")
        print(f"role_arn = arn:aws:iam::{account['account_id']}:role/admin")
        print()


def main():
    parser = argparse.ArgumentParser(
        description="Generate AWS CLI config for all accounts"
    )
    parser.add_argument(
        "--source-profile",
        default="default",
        help="Source profile name for AWS CLI config (default: default)",
    )

    args = parser.parse_args()

    aws_accounts = get_aws_accounts()
    generate_aws_cli_config(aws_accounts, args.source_profile)


def get_aws_accounts():
    aws_accounts = []

    for page in ssm_client.get_paginator('describe_parameters').paginate(
            Shared=True,
            ParameterFilters=[
                {
                    'Key': 'Name',
                    'Option': 'BeginsWith',
                    'Values': [
                        '/shared/aws-account-ids/',
                    ]
                },
            ]
    ):
        get_parameters_response = ssm_client.get_parameters(
            Names=[param['ARN'] for param in page['Parameters']],
        )
        aws_accounts.extend([json.loads(parameter['Value']) for parameter in get_parameters_response['Parameters']])

    return aws_accounts


if __name__ == '__main__':
    main()