import json
from google.cloud import bigquery
import os

table_id ="pingcap-gardener.aws_cost_data.dws_cost_aws_kubecost_asset_snapshot_di"

def dict_2_list(dict_data):
    """
    cc /collector/collect_kubecost_asset_snapshot.py
    The dict struct is @namespace_template
    Returns:
        list corresponding to _create_table_ function.
    """
    list_data = []
    for ns in dict_data:
        resource = {
            "namespace": ns["name"],
            "cluster": ns["cluster"],
            "minutes" : ns["metrics"]["minutes"],
            "cpuCores" : ns["metrics"]["cpuCores"],
            "cpuCoreRequestAverage" : ns["metrics"]["cpuCoreRequestAverage"],
            "cpuCoreUsageAverage" : ns["metrics"]["cpuCoreUsageAverage"],
            "cpuCoreHours" : ns["metrics"]["cpuCoreHours"],
            "cpuCost" : ns["metrics"]["cpuCost"],
            "cpuCostAdjustment" : ns["metrics"]["cpuCostAdjustment"],
            "cpuEfficiency" : ns["metrics"]["cpuEfficiency"],
            "gpuCount" : ns["metrics"]["gpuCount"],
            "gpuHours" : ns["metrics"]["gpuHours"],
            "gpuCost" : ns["metrics"]["gpuCost"],
            "gpuCostAdjustment" : ns["metrics"]["gpuCostAdjustment"],
            "networkTransferBytes" : ns["metrics"]["networkTransferBytes"],
            "networkReceiveBytes" : ns["metrics"]["networkReceiveBytes"],
            "networkCost" : ns["metrics"]["networkCost"],
            "networkCostAdjustment" : ns["metrics"]["networkCostAdjustment"],
            "loadBalancerCost" : ns["metrics"]["loadBalancerCost"],
            "loadBalancerCostAdjustment" : ns["metrics"]["loadBalancerCostAdjustment"],
            "pvBytes" : ns["metrics"]["pvBytes"],
            "pvByteHours" : ns["metrics"]["pvByteHours"],
            "pvCost" : ns["metrics"]["pvCost"],
            "pvCostAdjustment" : ns["metrics"]["pvCostAdjustment"],
            "ramBytes" : ns["metrics"]["ramBytes"],
            "ramByteRequestAverage" : ns["metrics"]["ramByteRequestAverage"],
            "ramByteUsageAverage" : ns["metrics"]["ramByteUsageAverage"],
            "ramByteHours" : ns["metrics"]["ramByteHours"],
            "ramCost" : ns["metrics"]["ramCost"],
            "ramCostAdjustment" : ns["metrics"]["ramCostAdjustment"],
            "ramEfficiency" : ns["metrics"]["ramEfficiency"],
            "sharedCost" : ns["metrics"]["sharedCost"],
            "externalCost" : ns["metrics"]["externalCost"],
            "totalCost" : ns["metrics"]["totalCost"],
            "totalEfficiency" : ns["metrics"]["totalEfficiency"],
            "date" : ns["date"],
        }
        list_data.append(resource)
    return list_data


def create_table():
    client = bigquery.Client()
    schema = [
        bigquery.SchemaField("namespace",  "STRING", mode="NULLABLE"),
        bigquery.SchemaField("cluster", "STRING", mode="NULLABLE"),
        bigquery.SchemaField("minutes", "FLOAT", mode="NULLABLE"),
        bigquery.SchemaField("cpuCores", "FLOAT", mode="NULLABLE"),
        bigquery.SchemaField("cpuCoreRequestAverage", "FLOAT", mode="NULLABLE"),
        bigquery.SchemaField("cpuCoreUsageAverage", "FLOAT", mode="NULLABLE"),
        bigquery.SchemaField("cpuCoreHours", "FLOAT", mode="NULLABLE"),
        bigquery.SchemaField("cpuCost", "FLOAT", mode="NULLABLE"),
        bigquery.SchemaField("cpuCostAdjustment", "FLOAT", mode="NULLABLE"),
        bigquery.SchemaField("cpuEfficiency", "FLOAT", mode="NULLABLE"),
        bigquery.SchemaField("gpuCount", "FLOAT", mode="NULLABLE"),
        bigquery.SchemaField("gpuHours", "FLOAT", mode="NULLABLE"),
        bigquery.SchemaField("gpuCost", "FLOAT", mode="NULLABLE"),
        bigquery.SchemaField("gpuCostAdjustment", "FLOAT", mode="NULLABLE"),
        bigquery.SchemaField("networkTransferBytes", "FLOAT", mode="NULLABLE"),
        bigquery.SchemaField("networkReceiveBytes", "FLOAT", mode="NULLABLE"),
        bigquery.SchemaField("networkCost", "FLOAT", mode="NULLABLE"),
        bigquery.SchemaField("networkCostAdjustment", "FLOAT", mode="NULLABLE"),
        bigquery.SchemaField("loadBalancerCost", "FLOAT", mode="NULLABLE"),
        bigquery.SchemaField("loadBalancerCostAdjustment", "FLOAT", mode="NULLABLE"),
        bigquery.SchemaField("pvBytes", "FLOAT", mode="NULLABLE"),
        bigquery.SchemaField("pvByteHours", "FLOAT", mode="NULLABLE"),
        bigquery.SchemaField("pvCost", "FLOAT", mode="NULLABLE"),
        bigquery.SchemaField("pvCostAdjustment", "FLOAT", mode="NULLABLE"),
        bigquery.SchemaField("ramBytes", "FLOAT", mode="NULLABLE"),
        bigquery.SchemaField("ramByteRequestAverage", "FLOAT", mode="NULLABLE"),
        bigquery.SchemaField("ramByteUsageAverage", "FLOAT", mode="NULLABLE"),
        bigquery.SchemaField("ramByteHours", "FLOAT", mode="NULLABLE"),
        bigquery.SchemaField("ramCost", "FLOAT", mode="NULLABLE"),
        bigquery.SchemaField("ramCostAdjustment", "FLOAT", mode="NULLABLE"),
        bigquery.SchemaField("ramEfficiency", "FLOAT", mode="NULLABLE"),
        bigquery.SchemaField("sharedCost", "FLOAT", mode="NULLABLE"),
        bigquery.SchemaField("externalCost", "FLOAT", mode="NULLABLE"),
        bigquery.SchemaField("totalCost", "FLOAT", mode="NULLABLE"),
        bigquery.SchemaField("totalEfficiency", "FLOAT", mode="NULLABLE"),
        bigquery.SchemaField("date", "DATE", mode="NULLABLE")
    ]
    # table_id = "pingcap-gardener.aws_cost_data_us_west1.aws_dev_resources_tree"
    table = bigquery.Table(table_id, schema=schema)
    table.time_partitioning = bigquery.TimePartitioning(
        type_=bigquery.TimePartitioningType.DAY,
        field="date",  # name of column to use for partitioning
    )
    table = client.create_table(table, exists_ok=True)  # Make an API request.

def overwrite_result_to_bq(data):
    _delete_table_partition(table_id, data)
    client = bigquery.Client()
    errors = client.insert_rows_json(table_id, data)  # Make an API request.
    if errors == []:
        print("New rows have been added.")
    else:
        print("Encountered errors while inserting rows: {}".format(errors))

def _delete_table_partition(table_id, data):
    partitions = []
    for row in data:
        partitions.append(row["date"])
    partitions = set(partitions)
    for date in partitions:
        date = date.replace("-", "")
        table_partition = "{}${}".format(table_id, date)
        print("Deleting table partition: {}".format(table_partition))
        client = bigquery.Client()
        errors = client.delete_table(table_partition, not_found_ok=True)
        if not errors :
            print("Partition {} is deleted. \n".format(table_partition))
        else:
            print("Encountered errors while deleting partition: {} \n".format(errors))


if __name__ == "__main__":
    resources_dict = json.load(open("kubecost_namespace_daily_asset_snapshot.json", 'r'))
    # TODO: ignore it when the shoot doesn't exist
    if os.path.exists('kubecost_namespace_daily_asset_snapshot.json'):
        resource_list = dict_2_list(resources_dict)
    else:
        resource_list = {}
    create_table()
    overwrite_result_to_bq(resource_list)