#!/bin/env python3
import subprocess
import datetime
import json
import calendar
import db_connect

ACCOUNT_LIST = ['pingcap','pingcap-dbaas','pingcap-test']

def get_aws_cost(account_name, start_date, end_date):
    cmd = 'aws ce get-cost-and-usage --time-period Start=' + start_date + ',End=' + end_date + ' --granularity MONTHLY --metrics "BlendedCost" "UnblendedCost" "UsageQuantity" --profile ' + account_name
    result = subprocess.check_output(cmd, shell=True)
    j = json.loads(result)
    return j['ResultsByTime'][0]['Total']['UnblendedCost']['Amount']

# Input month: 2020-12
def list_date_of_month(yyyy_mm):
    yyyy,mm = yyyy_mm.split("-")
    yyyy,mm = int(yyyy), int(mm)
    month_len = calendar.monthrange(yyyy,mm)[1] - 1
    first_day = datetime.date(yyyy,mm,1)
    date_list = [str(first_day)]
    next_day = first_day + datetime.timedelta(days=1)
    for i in range(month_len):
        date_list.append(str(next_day))
        next_day = next_day + datetime.timedelta(days=1)

    return date_list

# Input month: 2020-12
def insert_monthly_cost(account_name,yyyy_mm):
    monthly_dates = list_date_of_month(yyyy_mm)
    for date_str in monthly_dates:
        yyyy,mm,dd = date_str.split("-")
        this_date = datetime.date(int(yyyy),int(mm),int(dd))
        next_date_str = str(this_date + datetime.timedelta(days=1))
        # print(date_str + " " + next_date_str)
        daily_cost = get_aws_cost(account_name, date_str, next_date_str)
        print(account_name + " " + date_str + " " + daily_cost)
        insert_to_db(account_name,date_str,daily_cost)


# Input daily
def insert_daily_cost(account_name,yyyy_mm_dd):
    yyyy,mm,dd = yyyy_mm_dd.split("-")
    this_date = datetime.date(int(yyyy),int(mm),int(dd))
    next_date_str = str(this_date + datetime.timedelta(days=1))
    daily_cost = get_aws_cost(account_name, yyyy_mm_dd, next_date_str)
    print(account_name + " " + yyyy_mm_dd + " " + daily_cost)
    insert_to_db(account_name,yyyy_mm_dd,daily_cost)


# CREATE TABLE aws (account VARCHAR(255), date DATETIME, cost FLOAT(6,2));
# insert_to_db("pingcap","2020-12-01","3012.128")
def insert_to_db(account_name, date, cost):
    string_date_format = '%Y-%m-%d'
    parsed_datetime = datetime.datetime.strptime(date,string_date_format)
    sql = "INSERT INTO aws (account, date, cost) VALUES (%s, %s, %s)"
    val = (account_name, parsed_datetime, float(cost))
    db_connect.mycursor.execute(sql, val)
    db_connect.mydb.commit()
    

if __name__ == "__main__":
    for account in ACCOUNT_LIST:
        ## Batch update
        # insert_monthly_cost(account,"2021-01")

        ## Daily update
        current_date = datetime.datetime.today()
        day_before_yesterday = current_date - datetime.timedelta(2)
        day_before_yesterday = day_before_yesterday.strftime("%Y-%m-%d")
        insert_daily_cost(account,day_before_yesterday)