from io import StringIO
import json
import os
import re
import boto3
import datetime
import uuid
import csv
from datetime import timedelta
from google.auth import aws
from google.cloud import bigquery

billing_table = 'pingcap-gardener.dbaas_billing.gcp_billing_export_v1_010870_592F9D_20D797'
project = "pingcap-gardener"
env = {"PRODUCT": "'dbaas-prod'", "DEV": "'pingcap-gardener'", "OTHER": "('pingcap-gardener', 'dbaas-prod')"}
project_id = 'dbaas-prod'

platform = "gcp"
dimension = [
    'tenant',
    'component',
    'region',
    'service',
    'cost_amount',
    'cost_unit',
    'usage_amount',
    'usage_unit',
    'invoice_month',
    'cluster',
    'serviceType'
]

def query_gcp_cost(start_day, end_day, service_mapping, client, datasource, op):
    res = []

    query = (
        """
        SELECT
        IFNULL((select value from UNNEST(labels) where key='tenant'), "") as tenant,
        IFNULL((select value from UNNEST(labels) where key='cluster'), "") as cluster,
        IFNULL((select value from UNNEST(labels) where key='component'), "") as component,
        IFNULL((select value from UNNEST(labels) where key='servicetype'), "") as serviceType,
        location.location as region,
        sku.description as service,
        currency as cost_unit,
        usage.unit as usage_unit,
        invoice.month AS invoice_month,
        SUM(cost) as cost_amount,
        SUM(usage.amount) as usage_amount,
        SUM(IFNULL((SELECT SUM(c.amount) FROM UNNEST(credits) c), 0)) as total_credits
        FROM `pingcap-gardener.dbaas_billing.gcp_billing_export_v1_010870_592F9D_20D797`
        WHERE project.id {op} {project} AND usage_start_time >= '{start_time}' AND usage_start_time < '{end_time}'
        GROUP BY 1, 2, 3, 4, 5, 6, 7, 8, 9
        """.format(op=op, project=env[datasource], start_time=start_day, end_time=end_day)
    )
    query_job = client.query(query)
    results = query_job.result()  # Waits for job to complete.

    for row in results:
        r = {}
        for dim in dimension:
            if dim == "cost_amount":
                r[dim] = row[dim] + row["total_credits"]
            elif dim == "component":
                reresult=re.findall(r".-(.+?)$", row[dim])
                if len(reresult) <= 0:
                    r[dim] = row[dim]
                else:
                    r[dim] = reresult[0]
            elif dim == "service":
                r[dim] = row[dim]
                r["product"] = service_mapping.get(row[dim], "")
            else:
                r[dim] = row[dim]
        r["platform"] = platform
        r["date"] = str(start_day)
        r['environment'] = datasource
        res.append(r)
    
    return res

def bigquery_client():
    json_config_info = json.loads(os.getenv("GOOGLE_APPLICATION_CREDENTIALS"))
    credentials = aws.Credentials.from_info(json_config_info)
    scoped_credentials = credentials.with_scopes(
        ['https://www.googleapis.com/auth/cloud-platform'])
    
    client = bigquery.Client(project=project, credentials=scoped_credentials)
    return client

def load_product_mapping():
    service_mapping = {}
    csvfile = open("product_mapping.csv", "r")
    csvreader = csv.reader(csvfile)
    for item in csvreader:
        if csvreader.line_num == 1 or item[0] == '':
            continue
        service_mapping[item[1]] = item[0]

    return service_mapping

def LoadHistory(event, context):

    start_day = datetime.datetime.strptime(os.getenv("START_DAY"), '%Y-%m-%d').date()
    end_day = datetime.datetime.strptime(os.getenv("END_DAY"), '%Y-%m-%d').date()

    bigquery = bigquery_client()

    sqs = boto3.client("sqs")
    queue = sqs.get_queue_url(QueueName=os.getenv("QUEUE_NAME"))

    service_mapping = load_product_mapping()

    while start_day <= end_day: 
        next_day = start_day + timedelta(1)
        reports = query_gcp_cost(start_day, next_day, service_mapping, bigquery, os.getenv('ENV'), os.getenv('op'))

        for r in reports:
            sqs.send_message(
                QueueUrl=queue['QueueUrl'], 
                MessageBody=json.dumps(r),
                MessageDeduplicationId=uuid.uuid4().hex,
                MessageGroupId=os.getenv("MESSAGE_GROUP_ID")
            )  
        start_day = start_day + timedelta(1)

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

def lambda_handler(event, context):
    today=datetime.date.today() - timedelta(2)
    yesterday = today - timedelta(1)
    service_mapping = load_product_mapping()
    bigquery = bigquery_client()
    reports = query_gcp_cost(yesterday, today, service_mapping, bigquery, "PRODUCT", '=')
    reports.extend(query_gcp_cost(yesterday, today, service_mapping, bigquery, "DEV", '='))
    reports.extend(query_gcp_cost(yesterday, today, service_mapping, bigquery, "OTHER", 'NOT IN'))
    
    sqs = boto3.client("sqs")
    queue = sqs.get_queue_url(QueueName=os.getenv("QUEUE_NAME"))

    for r in reports:
        sqs.send_message(
            QueueUrl=queue['QueueUrl'], 
            MessageBody=json.dumps(r),
            MessageDeduplicationId=uuid.uuid4().hex,
            MessageGroupId=os.getenv("MESSAGE_GROUP_ID")
        ) 

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