import boto3
import uuid
import json
from datetime import datetime

now = datetime.now()
year = now.strftime("%Y")
month = now.strftime("%m")
day = now.strftime("%d")
file_name = 'query-data.json'

# Finalized data
# Appends finalized data to the last day of the previous month in S3

# last_month = now.replace(day=1) - timedelta(days=1)
# year = last_month.strftime("%Y")
# month = last_month.strftime("%m")
# day = last_month.strftime("%d")
# file_name = 'finalized-data.json'

# Vars to Change!
integration_uuid = "CHANGEME"                                    # Cost Management integration_uuid
bucket = 'CHANGEME'                                         # Bucket created for query results
database = 'CHAGNEME'                                       # Database to execute athena queries
export_name = 'CHANGEME'                                    # Name of your data export for athena
output=f's3://{bucket}/{year}/{month}/{day}/{uuid.uuid4()}' # Output location for query results

# Athena query
query = f"SELECT * FROM {database}.{export_name} WHERE ((bill_billing_entity = 'AWS Marketplace' AND line_item_legal_entity like '%Red Hat%') OR (line_item_legal_entity like '%Amazon Web Services%' AND line_item_line_item_description like '%Red Hat%') OR (line_item_legal_entity like '%Amazon Web Services%' AND line_item_line_item_description like '%RHEL%') OR (line_item_legal_entity like '%AWS%' AND line_item_line_item_description like '%Red Hat%') OR (line_item_legal_entity like '%AWS%' AND line_item_line_item_description like '%RHEL%') OR (line_item_legal_entity like '%AWS%' AND product_product_name like '%Red Hat%') OR (line_item_legal_entity like '%Amazon Web Services%' AND product_product_name like '%Red Hat%')) AND year = '{year}' AND month = '{month}'"

def lambda_handler(event, context):
    # Initiate Boto3 athena Client
    athena_client = boto3.client('athena')

    # Trigger athena query
    response = athena_client.start_query_execution(
        QueryString=query,
        QueryExecutionContext={
            'Database': database
        },
        ResultConfiguration={
            'OutputLocation': output
        }
    )

    # Save query execution to s3 object
    s3 = boto3.client('s3')
    json_object = {"integration_uuid": integration_uuid, "bill_year": year, "bill_month": month, "query_execution_id": response.get("QueryExecutionId"), "result_prefix": output}
    s3.put_object(
        Body=json.dumps(json_object),
        Bucket=bucket,
        Key=file_name
    )

    return json_object