import csv
import json
from datetime import timedelta, datetime, timezone

import boto3
from dateutil import parser

ONE_DAY = 60 * 60 * 24

SAVINGS_PLAN_FIELDS = [
    'accountId',
    'accountName',
    'id',
    'description',
    'commitment',
    'paymentOption',
    'state',
    'start',
    'end',
]

RI_FIELDS = [
    'accountId',
    'accountName',
    'service',
    'id',
    'instanceClass',
    'count',
    'paymentOption',
    'state',
    'start',
    'end',
]

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

def main():
    current_user = iam_client.get_user()
    username = current_user['User']['UserName']
    aws_accounts = get_aws_accounts()

    generate_savings_plans_report(aws_accounts, username)
    generate_ris_report(aws_accounts, username)


def generate_savings_plans_report(aws_accounts, username):
    with open('savingsplans.csv', 'w') as csvfile:
        writer = csv.DictWriter(csvfile, fieldnames=SAVINGS_PLAN_FIELDS)
        writer.writeheader()

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

                savingsplans_client = get_client(username, account, 'savingsplans')

                for savings_plan in savingsplans_client.describe_savings_plans(
                        maxResults=1000
                )['savingsPlans']:
                    start = parser.parse(savings_plan['start'])
                    end = parser.parse(savings_plan['end'])

                    if end < datetime.now(timezone.utc) - timedelta(days=90):
                        print(f'Not including savings plan {savings_plan["savingsPlanId"]} as it expired more than 90 days ago')
                        continue

                    row = {
                        'accountId': account['account_id'],
                        'accountName': account['name'],
                        'id': savings_plan['savingsPlanId'],
                        'description': savings_plan['description'],
                        'commitment': savings_plan['commitment'],
                        'paymentOption': savings_plan['paymentOption'],
                        'state': savings_plan['state'],
                        'start': start.strftime('%Y-%m-%d'),
                        'end': end.strftime('%Y-%m-%d'),
                    }
                    writer.writerow(row)
            except Exception as e:
                print(f'Error gathering savings plan data for account {account["name"]}: {e}')


def generate_ris_report(aws_accounts, username):
    with open('ris.csv', 'w') as csvfile:
        writer = csv.DictWriter(csvfile, fieldnames=RI_FIELDS)
        writer.writeheader()

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

                write_rds_ris(account, username, writer)
                write_elasticache_ris(account, username, writer)
                write_es_ris(account, username, writer)
            except Exception as e:
                print(f'Error gathering RI data for account {account["name"]}: {e}')


def write_es_ris(account, username, writer):
    es_client = get_client(username, account, 'es')
    for page in es_client.get_paginator(
            'describe_reserved_elasticsearch_instances').paginate():
        for ri in page['ReservedElasticsearchInstances']:
            start = ri['StartTime']
            end = start + timedelta(seconds=ri['Duration'])

            if end < datetime.now(timezone.utc) - timedelta(days=90):
                print(f'Not including ES RI {ri["ReservedElasticsearchInstanceId"]} as it expired more than 90 days ago')
                continue

            row = {
                'accountId': account['account_id'],
                'accountName': account['name'],
                'service': 'Elasticsearch',
                'id': ri['ReservedElasticsearchInstanceId'],
                'instanceClass': ri['ElasticsearchInstanceType'],
                'count': ri['ElasticsearchInstanceCount'],
                'paymentOption': ri['PaymentOption'],
                'state': ri['State'],
                'start': start.strftime('%Y-%m-%d'),
                'end': end.strftime('%Y-%m-%d')
            }
            writer.writerow(row)


def write_elasticache_ris(account, username, writer):
    elasticache_client = get_client(username, account, 'elasticache')
    for page in elasticache_client.get_paginator(
            'describe_reserved_cache_nodes').paginate():
        for ri in page['ReservedCacheNodes']:
            start = ri['StartTime']
            end = start + timedelta(seconds=ri['Duration'])

            if end < datetime.now(timezone.utc) - timedelta(days=90):
                print(f'Not including Elasticache RI {ri["ReservedCacheNodeId"]} as it expired more than 90 days ago')
                continue

            row = {
                'accountId': account['account_id'],
                'accountName': account['name'],
                'service': 'Elasticache',
                'id': ri['ReservedCacheNodeId'],
                'instanceClass': ri['CacheNodeType'],
                'count': ri['CacheNodeCount'],
                'paymentOption': ri['OfferingType'],
                'state': ri['State'],
                'start': start.strftime('%Y-%m-%d'),
                'end': end.strftime('%Y-%m-%d')
            }
            writer.writerow(row)


def write_rds_ris(account, username, writer):
    rds_client = get_client(username, account, 'rds')
    for page in rds_client.get_paginator(
            'describe_reserved_db_instances').paginate():
        for ri in page['ReservedDBInstances']:
            start = ri['StartTime']
            end = start + timedelta(seconds=ri['Duration'])

            if end < datetime.now(timezone.utc) - timedelta(days=90):
                print(f'Not including RDS RI {ri["ReservedDBInstanceId"]} as it expired more than 90 days ago')
                continue

            row = {
                'accountId': account['account_id'],
                'accountName': account['name'],
                'service': 'RDS',
                'id': ri['ReservedDBInstanceId'],
                'instanceClass': ri['DBInstanceClass'],
                'count': ri['DBInstanceCount'],
                'paymentOption': ri['OfferingType'],
                'state': ri['State'],
                'start': start.strftime('%Y-%m-%d'),
                'end': end.strftime('%Y-%m-%d')
            }
            writer.writerow(row)


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


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

        credentials = assume_role_response["Credentials"]
        return boto3.client(service,
            aws_access_key_id=credentials["AccessKeyId"],
            aws_secret_access_key=credentials["SecretAccessKey"],
            aws_session_token=credentials["SessionToken"],
        )
    else:
        return boto3.client(service)


if __name__ == '__main__':
    main()