import time
import boto3
import os
import json
import uuid
import datetime

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():
    athenaClient = boto3.client('athena')
    query = '''
    SELECT
        to_iso8601(current_date - interval '2' day) 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_blended_cost ) as cost_amount,
        pricing_currency as cost_unit,
        line_item_usage_type as service
    FROM "athenacurcfn_prod_billing"."prod_billing"
    where
        to_iso8601(line_item_usage_start_date) <= to_iso8601(current_date - interval '1' day)
        and to_iso8601(line_item_usage_start_date)  >=  to_iso8601(current_date - interval '2' day)
    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
    '''
    response = athenaClient.start_query_execution(
        QueryString=query,
        QueryExecutionContext={
            'Database': "athenacurcfn_prod_billing",
        },
        ResultConfiguration = {
            'OutputLocation': 's3://prod-billing-reports-athena/tempResult/',
        }
    )

    # print(response)
    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


def lambda_handler(event, context):
    result = fetch_aws_costs()

    sqs = boto3.client('sqs')
    queue = sqs.get_queue_url(QueueName=os.getenv("QUEUE_NAME"))
    for r in result:
        response = sqs.send_message(
            QueueUrl=queue['QueueUrl'],
            MessageBody=json.dumps(r),
            MessageDeduplicationId=uuid.uuid4().hex,
            MessageGroupId=os.getenv("MESSAGE_GROUP_ID")
        )
        print( response )

    return {
        'statusCode': 200,
        'body': json.dumps('ok')
    }