import datetime
import sys
import time
import boto3
import os
import json
import uuid
import csv

cost_analysis_dicts = {
    "Amazon Elastic File System": "EFS",
    "AmazonCloudWatch": "CloudWatch",
    "AWS Cost Explorer": "Cost Explorer",
    "Elastic Load Balancing": "ELB",
    "AWS CloudTrail": "CloudTrail",
    "AWS Data Transfer": "DataTransfer",
    "Amazon Elastic Compute Cloud": "EC2",
    "Amazon Simple Notification Service": "SNS",
    "Amazon Athena": "Athena",
    "Amazon Simple Storage Service": "S3",
    "AWS Key Management Service": "KMS",
    "Amazon Simple Queue Service": "SQS",
    "Amazon EC2 Container Registry (ECR)": "ECR",
    "Amazon Relational Database Service": "RDS",
    "AWS WAF": "WAF",
    "Amazon Route 53": 'Route 53',
    "Amazon S3 Glacier Deep Archive": "S3",
    "Amazon CloudFront": "CloudFront",
    "AWS Certificate Manager": 'Certificate Manager',
    "AWS Lambda": "Lambda",
    "Savings Plans for AWS Compute usage": "SavingPlan",
    "Amazon Elastic Container Service for Kubernetes": "EKS",
    "AWS Global Accelerator": "Global Accelerator"
}

def fetch_aws_costs(begin):
    print(begin)
    beginDate = datetime.datetime.strptime(begin, '%Y-%m-%d').date()
    athenaClient = boto3.client('athena')
    query = '''
    SELECT
        to_iso8601(date '{}') as date, 
        resource_tags_user_tenant as tenant,
        resource_tags_user_cluster as cluster,
        resource_tags_user_component as component,
        product_product_name as product,
        product_region as region,
        resource_tags_user_servicetype as servicetype,
        sum( line_item_unblended_cost ) as cost_amount,
        pricing_currency as cost_unit,
        line_item_usage_type as service
    FROM "athenacurcfn_dev_billing"."dev_billing"
    where
        to_iso8601(line_item_usage_start_date) <= to_iso8601(date '{}')
        and to_iso8601(line_item_usage_start_date)  >=  to_iso8601(date '{}')
    group by
        resource_tags_user_tenant,
        resource_tags_user_cluster,
        resource_tags_user_component,
        line_item_usage_type,
        product_region,
        product_product_name,
        resource_tags_user_servicetype,
        pricing_currency
    '''.format(beginDate,
               (beginDate + datetime.timedelta(days=1)) ,
               beginDate
               )

    response = athenaClient.start_query_execution(
        QueryString=query,
        QueryExecutionContext={
            'Database': "athenacurcfn_dev_billing",
        },
    )

    queryExecutionId = response['QueryExecutionId']

    while True:
        response = athenaClient.batch_get_query_execution(
            QueryExecutionIds = [
                queryExecutionId
            ],
        )

        if len (response['QueryExecutions']) < 1:
            time.sleep(2)
            continue

        queryResult = response['QueryExecutions'][0]

        queryState = queryResult['Status']['State']
        if queryState in ['QUEUED', 'RUNNING']:
            print("querying!")
            time.sleep(2)
        elif queryState in ['FAILED', 'CANCELLED']:
            print(response)
            raise Exception("Athena Query Failed")
        else:
            break

    response = athenaClient.get_query_results(
        QueryExecutionId=queryExecutionId,
        MaxResults=1000
    )
    result = []
    while True:
        # print( response )
        if 'ResultSet' not in response:
            return result
        if 'Rows' not in response['ResultSet']:
            return result

        if len( response['ResultSet']['Rows']) <= 9:
            return result
        for row in response['ResultSet']['Rows'][1:]:
            rowData = row['Data']
            serviceType = rowData[4]['VarCharValue']
            # for cost_analysis in cost_analysis_dicts:
            #     for condition in cost_analysis[0]:
            #         if condition in rowData[4]['VarCharValue']:
            #             serviceType = cost_analysis[1]
            if serviceType in cost_analysis_dicts:
                serviceType = cost_analysis_dicts[serviceType]

            result.append(
                {
                    "date": rowData[0]['VarCharValue'],
                    "tenant": rowData[1]['VarCharValue'],
                    "cluster": rowData[2]['VarCharValue'],
                    "component": rowData[3]['VarCharValue'],
                    "product": serviceType,
                    "region": rowData[5]['VarCharValue'],
                    "serviceType": rowData[6]['VarCharValue'],
                    "cost_amount": float( rowData[7]['VarCharValue'] ),
                    "cost_unit": rowData[8]['VarCharValue'],
                    "service": rowData[9]['VarCharValue'],
                    "platform": "aws"
                }
            )

        if 'NextToken' not in response:
            break

        nextToken = response['NextToken']
        print( nextToken )
        response = athenaClient.get_query_results(
            QueryExecutionId=queryExecutionId,
            MaxResults=1000,
            NextToken= nextToken,
        )

    return result


csvFile = open(sys.argv[1] + '-costs.csv', 'w+')
fieldnames = ['date', 'tenant', 'cluster',
              'component', 'product', 'region',
              'serviceType', "cost_amount",
              'cost_unit', 'service', "platform"]
result = fetch_aws_costs(sys.argv[1])
writer = csv.DictWriter(csvFile, fieldnames=fieldnames)
writer.writeheader()
for r in result:
    writer.writerow({
        "date": r['date'],
        "tenant": r['tenant'],
        "cluster":  r["cluster"],
        "component": r['component'],
        'product': r['product'],
        'region': r['region'],
        "serviceType": r['serviceType'],
        "cost_amount": r['cost_amount'],
        'cost_unit': r['cost_unit'],
        'service': r['service'],
        'platform': r['platform']
    })

csvFile.close()