import csv
import json

import boto3

FIELDS = [
    'accountId',
    'accountName',
    'actionType',
    'currencyCode',
    'currentResourceSummary',
    'currentResourceType',
    'estimatedMonthlyCost',
    'estimatedMonthlySavings',
    'estimatedSavingsPercentage',
    'implementationEffort',
    'lastRefreshTimestamp',
    'recommendationId',
    'recommendationLookbackPeriodInDays',
    'recommendedResourceSummary',
    'recommendedResourceType',
    'region',
    'resourceArn',
    'resourceId',
    'restartNeeded',
    'rollbackPossible',
    'source',
    'tags',
]

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

current_user = iam_client.get_user()
user_name = current_user['User']['UserName']

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']])

with open('recommendations.csv', 'w') as csvfile:
    writer = csv.DictWriter(csvfile, fieldnames=FIELDS)
    writer.writeheader()

    for account in aws_accounts:
        try:
            print(f'Gathering data for account {account["name"]}')

            if not account['name'] == 'prod':
                print(f'Assuming role {account["role"]}')
                assume_role_response = sts_client.assume_role(
                    RoleArn=f"arn:aws:iam::{account['account_id']}:role/{account['role']}",
                    RoleSessionName=user_name
                )

                credentials = assume_role_response["Credentials"]
                cost_optimization_hub_client = boto3.client('cost-optimization-hub',
                    aws_access_key_id=credentials["AccessKeyId"],
                    aws_secret_access_key=credentials["SecretAccessKey"],
                    aws_session_token=credentials["SessionToken"],
                )
            else:
                cost_optimization_hub_client = boto3.client('cost-optimization-hub')

            print(f'Retrieving data from cost optimization hub')
            for page in cost_optimization_hub_client.get_paginator('list_recommendations').paginate():
                for item in page['items']:
                    item['accountName'] = account['name']
                    writer.writerow(item)
        except Exception as e:
            print(f'Error gathering RI data for account {account["name"]}: {e}')