''' Create report from Cost exporer API.
In manual mode by default creates a report for the last 12 months.
In auto mode creates two reports for last and current month and uploads them to s3. Also updates quicksight
'''
import pandas as pd
import time
import sys
import boto3
import argparse

from datetime import date
from datetime import timedelta

import core
import accounts


MIN_MONTHLY_REPORT_DAYS = 1
AWS_REGION = "us-east-1"

QS_S3_DATA_SET = "Billing reports s3"
QS_ANALYSIS = "Analysis - Fully Tagged Costs"
QS_DASHBOARD = "Amortized - Fully Tagged Costs"

def parse_arguments():
    parser = argparse.ArgumentParser()

    parser.add_argument('--manual', action='store_true', help='generate report from local Cost Explorer of AWS accounts')
    parser.add_argument('--start-date', default='', help='From which date generate reports (only for --manual)')
    parser.add_argument('--end-date', default='', help='To which date generate reports (only for --manual)')
    parser.add_argument('--out-file', default='', help='Name of the output file (only for --manual)')

    parser.add_argument('--auto', action='store_true', help='generate report from local Cost Explorer of AWS accounts and upload to s3')
    parser.add_argument('--bucket', default='prod-infra-billing-reports', help='S3 bucket where report will be uploaded')
    parser.add_argument('--prefix', default='', help='prefix to add to uploaded s3 objects')

    return parser.parse_args()

# Get dataframe based on Cost Explorer report grouped by tag plat_env_project_service
# for given dates
def getTaggedReportDf(client, startDate: str, endDate: str) -> pd.DataFrame:
    results = []
    results_by_day={}
    # Request parameters
    timePeriod = {
        'Start': startDate,
        'End': endDate
    }
    groupBy = [{
        'Type': 'TAG',
        'Key': 'plat_env_project_service'
    }]
    granularity = 'DAILY'
    filter = {'Not': {'Dimensions': { 'Key': 'RECORD_TYPE', 'Values': ['Credit'] }}}
    metrics = ['NetAmortizedCost']

    response = client.get_cost_and_usage(
        TimePeriod=timePeriod,
        GroupBy=groupBy,
        Granularity=granularity,
        Filter=filter,
        Metrics=metrics,
    )
    results.extend(response['ResultsByTime'])
    while 'NextPageToken' in response:
        nextToken = response['NextPageToken']
        response = client.get_cost_and_usage(
            TimePeriod=timePeriod,
            GroupBy=groupBy,
            Granularity=granularity,
            Filter=filter,
            Metrics=metrics,
            NextPageToken=nextToken,
        )
        results.extend(response['ResultsByTime'])
        if 'NextPageToken' in response:
            nextToken = response['NextPageToken']
        else:
            nextToken = False

    for day_entry in results:
        date_key = day_entry['TimePeriod']['Start']
        if date_key not in results_by_day:
            results_by_day[date_key] = {'plat_env_project_service' : date_key}

        for group in day_entry['Groups']:
            tag = group['Keys'][0].replace('plat_env_project_service$', '')
            # Ignore aggregated data for a day
            if tag != '':
                results_by_day[date_key][tag.upper()] = float(group['Metrics']['NetAmortizedCost']['Amount'])

    out_df = pd.DataFrame(list(results_by_day.values()))
    out_df = out_df.fillna(0)
    out_df.set_index('plat_env_project_service', inplace = True)
    return out_df

# Get dataframe based on Cost Explorer report grouped by Service
# for all untagged resources
def getUntaggedReportDf(client, startDate: str, endDate: str) -> pd.DataFrame:
    results = []
    results_by_day={}
    # Request parameters
    timePeriod = {
        'Start': startDate,
        'End': endDate
    }
    groupBy = [{
        'Type': 'DIMENSION',
        'Key': 'SERVICE'
    }]
    granularity = 'DAILY'
    filter = {
        'And': [
            {'Not': {'Dimensions': { 'Key': 'RECORD_TYPE', 'Values': ['Credit'] }}},
            {'Tags': {'Key': 'plat_env_project_service', 'MatchOptions': ['ABSENT']}}
    ]}
    metrics = ['NetAmortizedCost']

    response = client.get_cost_and_usage(
        TimePeriod=timePeriod,
        GroupBy=groupBy,
        Granularity=granularity,
        Filter=filter,
        Metrics=metrics,
    )
    results.extend(response['ResultsByTime'])
    while 'NextPageToken' in response:
        nextToken = response['NextPageToken']
        response = client.get_cost_and_usage(
            TimePeriod=timePeriod,
            GroupBy=groupBy,
            Granularity=granularity,
            Filter=filter,
            Metrics=metrics,
            NextPageToken=nextToken,
        )
        results.extend(response['ResultsByTime'])
        if 'NextPageToken' in response:
            nextToken = response['NextPageToken']
        else:
            nextToken = False

    for day_entry in results:
        date_key = day_entry['TimePeriod']['Start']
        if date_key not in results_by_day:
            results_by_day[date_key] = {'Service' : date_key}

        for group in day_entry['Groups']:
            tag = group['Keys'][0]
            results_by_day[date_key][tag] = float(group['Metrics']['NetAmortizedCost']['Amount'])

    out_df = pd.DataFrame(list(results_by_day.values()))
    out_df = out_df.fillna(0)
    out_df.set_index('Service', inplace = True)
    return out_df

# Create report for given dates based on local CE reports from each account.
def createReport(outFile: str, startDate: str, endDate: str) -> pd.DataFrame:
    resultDf = core.createEmptyDf()
    for acc in accounts.accountList():
        session = boto3.Session(profile_name=acc['profile'])
        ce_client = session.client('ce')
        amortizedDf = core.generateAmortizedCosts(getTaggedReportDf(ce_client, startDate, endDate))
        accountDf = core.generateAccountCost(getUntaggedReportDf(ce_client, startDate, endDate), acc['platform'], acc['env'])
        resultDf = pd.concat([resultDf, amortizedDf, accountDf], ignore_index=True)
    resultDf = core.sortBillingDf(resultDf)

    summarizedDf = core.generateSummaries(resultDf)
    print(summarizedDf)
    summarizedDf.to_csv(outFile, index=False)
    return summarizedDf


def autoModeS3Reports(bucket: str, prefix='', awsProfile='gdb-infra-prod'):
    '''
    This function create reports for previous and current months
    and uploads them to s3. Current month is excluded if
    current day is less than MIN_MONTHLY_REPORT_DAYS
    '''
    session = boto3.Session(profile_name=awsProfile,region_name = AWS_REGION)
    s3 = session.resource('s3')

    today = date.today()
    currentMonthFirst = today.replace(day=1)
    previousMonthFirst = (currentMonthFirst - timedelta(days=1)).replace(day=1)
    currentMonth = currentMonthFirst.strftime("%Y-%m-%d")
    previousMonth = previousMonthFirst.strftime("%Y-%m-%d")

    # Create current month only if there were at least N days in it
    if today.day > MIN_MONTHLY_REPORT_DAYS:
        endDate = today.strftime("%Y-%m-%d")
        outFileCurrentMonth = "report-" + currentMonthFirst.strftime("%Y-%m") + ".csv"
        print([outFileCurrentMonth, currentMonth, endDate])
        createReport(outFileCurrentMonth, currentMonth, endDate)
        s3.Bucket(bucket).upload_file(outFileCurrentMonth, prefix + outFileCurrentMonth)

    # Previous month report
    outFilePreviousMonth = "report-" + previousMonthFirst.strftime("%Y-%m") + ".csv"
    print([outFilePreviousMonth, previousMonth, currentMonth])
    createReport(outFilePreviousMonth, previousMonth, currentMonth)
    s3.Bucket(bucket).upload_file(outFilePreviousMonth, prefix + outFilePreviousMonth)

def updateQuickSight(awsProfile='gdb-delphi-dev'):
    ''' Refreshes quicksight s3 data set, publishes latest dashboard
    '''
    session = boto3.Session(profile_name=awsProfile,region_name = AWS_REGION)
    client = session.client('quicksight')
    awsAccId = accounts.accIdByProfile(awsProfile)
    # Refresh S3 dataset
    datasetId = ''
    ingestionId = str(int(time.time()))

    response = client.list_data_sets(AwsAccountId=awsAccId)
    for ds in response['DataSetSummaries']:
        if ds['Name'] == QS_S3_DATA_SET:
            datasetId = ds['DataSetId']
            break

    response = client.create_ingestion(
        DataSetId=datasetId,
        IngestionId=ingestionId,
        AwsAccountId=awsAccId)
    while True:
        response = client.describe_ingestion(
            DataSetId=datasetId,
            IngestionId=ingestionId,
            AwsAccountId=awsAccId)
        if response['Ingestion']['IngestionStatus'] in ('INITIALIZED', 'QUEUED', 'RUNNING'):
            time.sleep(10) #change sleep time according to your dataset size
        elif response['Ingestion']['IngestionStatus'] == 'COMPLETED':
            print("refresh completed. RowsIngested {0}, RowsDropped {1}, IngestionTimeInSeconds {2}, IngestionSizeInBytes {3}".format(
                response['Ingestion']['RowInfo']['RowsIngested'],
                response['Ingestion']['RowInfo']['RowsDropped'],
                response['Ingestion']['IngestionTimeInSeconds'],
                response['Ingestion']['IngestionSizeInBytes']))
            break
        else:
            print("refresh failed! - status {0}".format(response['Ingestion']['IngestionStatus']))
            sys.exit(1)
    # Dashboard is updated automatically upon data set refresh

if __name__ == '__main__':
    args = parse_arguments()
    start = time.time()
    if not (args.manual or args.auto):
        print("Either --manual or --auto should be set")
        exit(1)

    if args.manual:
        #Default start day - 1 day of the month one year ago
        if args.start_date == '':
            start_date = date.today() - timedelta(days=364+date.today().day)
            args.start_date = start_date.strftime("%Y-%m-%d")
        #Default stop day - yesterday
        if args.end_date == '':
            args.end_date = (date.today() - timedelta(days=1)).strftime("%Y-%m-%d")
        #Default out_file name:
        if args.out_file == '':
            args.out_file = 'report-from-' + args.start_date + '-to-' + args.end_date + '.csv'
        print(args)
        createReport(args.out_file, args.start_date, args.end_date)

    if args.auto:
        autoModeS3Reports(args.bucket)
        updateQuickSight()


    end = time.time()
    print(end - start)