import datetime
import os
from datetime import datetime, timedelta
import concurrent.futures
import mysql.connector
from airflow import models
from airflow.operators.python import PythonOperator
from airflow.providers.google.cloud.operators import bigquery

YESTERDAY = datetime.now() - timedelta(days=1)
default_args = {
    'owner': 'zhuxiaopeng',
    'depends_on_past': False,
    'email': ['xiaopeng.zhu@pingcap.com'],
    'email_on_failure': False,
    'email_on_retry': False,
    'retries': 1,
    'retry_delay': timedelta(minutes=5),
    'start_date': YESTERDAY,
}

with models.DAG(
        dag_id='pch_billing',
        catchup=False,
        description="send pch billing data to mysql",
        default_args=default_args,
        schedule_interval="0 */12 * * *",
) as dag:
    billing_table = "`gcp_billing_v2.gcp_billing_export_v1_01A169_D3D9BF_FB82A2`"
    full_table = "`gcp_billing_v2.gcp_billing_export_v1_*`"


    def get_connection():
        conf = {
            'user': os.getenv("username"),
            'password': os.getenv("password"),
            'host': os.getenv("host"),
            'database': os.getenv("database"),
            'raise_on_warnings': True
        }


        return mysql.connector.connect(**conf)


    def pch_billing():
        after_date = start_date(False)
        bq_hook = bigquery.BigQueryHook(use_legacy_sql=False)

        with get_connection() as cnx:
            with cnx.cursor() as cursor:
                cursor.execute(f"DELETE FROM cloud_product_cost WHERE date >= '{after_date}'")
                cnx.commit()
        cli = bq_hook.get_client()

        sync_aws(after_date, cli)
        sync_gcp(after_date, cli)
        sync_azure(after_date, cli)


    def sync_gcp(after_date, cli):
        query_sql = contract_gcp_query(after_date)

        results = cli.query(query_sql).result()

        with get_connection() as cnx:
            with cnx.cursor() as cursor:
                # build map from cloudInsID to resourceID
                cursor.execute(
                    "select id,cloudInsID from cloud_resource where providerID = 2 and classType <> 'Prod' and isDel = 0")
                resource_dict = {}
                for (id, cloudInsID) in cursor.fetchall():
                    resource_dict[cloudInsID] = id
                cnx.commit()

        # deal gcp billing data
        gcp_billing = []

        count = 0
        for row in results:

            if row[3] is None or row[3] == "":
                continue
            if row[3] not in resource_dict:
                continue
            gcp_billing.append(
                {
                    "create_at": datetime.now(),
                    "resource_id": resource_dict[row[3]],
                    "product": row[2],
                    "cost": row[1],
                    "date": datetime(row[0].year, row[0].month, row[0].day)
                }
            )
            count += 1

        insert_db(gcp_billing)

        print(f"insert {count} gcp billing data")


    def insert_db(billing):
        with get_connection() as cnx:
            with cnx.cursor() as cursor:
                cnx.start_transaction()

                try:
                    data_to_insert = [
                        (item['create_at'], item['resource_id'], item['product'], item['cost'], item['date'])
                        for
                        item
                        in billing]
                    cursor.executemany(
                        "INSERT INTO cloud_product_cost (createdAt,resourceID,productName,cost,date) VALUES (%s,%s,%s,%s,%s)",
                        data_to_insert)
                    cnx.commit()
                except Exception as e:
                    cnx.rollback()
                    raise e


    def start_date(full):
        if full:
            return "2022-07-01"
        before45d = datetime.now() - timedelta(days=45)
        return before45d.strftime("%Y-%m-%d")


    def contract_gcp_query(sdate):
        table = billing_table
        if sdate <= "2023-04-19":
            table = full_table
        query_sql = "SELECT `date`, ROUND(SUM( (cost + credits_amount) * currency_conversion_rate ), 4) AS cost, service.description AS product, project_id AS projID FROM ( SELECT *  FROM ( SELECT sku.description AS product, DATE( usage_start_time ) AS date, project.id AS project_id, IF(( SELECT SUM(x) FROM UNNEST(ARRAY( SELECT amount FROM UNNEST(credits))) AS x) IS NULL, 0, ( SELECT SUM(x) FROM UNNEST(ARRAY( SELECT amount FROM UNNEST(credits))) AS x )) AS credits_amount, * FROM"
        query_sql += table
        query_sql += f"WHERE DATE(usage_start_time) >= '{sdate}' ) T )"
        query_sql += " GROUP BY projID, product, date HAVING cost > 0 ORDER BY projID ASC, date ASC, cost DESC"
        return query_sql


    def sync_aws(sdate, cli):
        with get_connection() as cnx:
            with cnx.cursor() as cursor:
                query = (
                    "SELECT id,cloudInsID FROM cloud_resource where providerID = 1 and classType <> 'Prod' and isDel = 0")
                cursor.execute(query)
                aws_accounts = []
                for (id, cloudInsID) in cursor.fetchall():
                    aws_accounts.append((id, cloudInsID))
                cnx.commit()

        with concurrent.futures.ThreadPoolExecutor(max_workers=10) as executor:
            future_results = []
            for (id, cloudInsID) in aws_accounts:
                future_results.append(executor.submit(get_aws_billing_by_id, id, cloudInsID, sdate, cli))
            for future in concurrent.futures.as_completed(future_results):
                try:
                    future.result()
                except Exception as e:
                    print(e)


    def get_aws_billing_by_id(resource_id, aws_id, sdate, cli):

        query = construct_aws_query(aws_id, sdate)

        sync_aws_billing(resource_id, aws_id, query, sdate, cli)


    def construct_aws_query(cid, sdate):
        table = f"aws_nonprod_billing_native.{cid}_aws_billing_data"
        day = "DATE(line_item_usage_start_date) as day"
        cost_measure = "ROUND(SUM(line_item_unblended_cost), 4) as cost"
        date_range = f"DATE(line_item_usage_start_date) >= '{sdate}'"

        query = f"(SELECT {day}, {cost_measure}, line_item_product_code as product FROM {table} WHERE "
        query += f"{date_range} AND (line_item_line_item_type <> 'SavingsPlanNegation' AND line_item_product_code <> 'ComputeSavingsPlans')"
        query += " GROUP BY day, product HAVING cost <> 0 ORDER BY day ASC, cost DESC) UNION ALL (SELECT "
        query += f"{day}, {cost_measure}, 'SavingsPlans' as product FROM {table} WHERE {date_range}"
        query += " AND (line_item_line_item_type = 'SavingsPlanNegation' OR line_item_product_code = 'ComputeSavingsPlans')"
        query += " GROUP BY day, product ORDER BY day ASC, cost DESC)"

        return query


    def sync_aws_billing(resource_id, aws_id, query, sdate, cli):
        res = cli.query(query).result()

        billing_data = []
        for row in res:
            date = row[0]
            cost = row[1]
            product = row[2]

            billing_data.append(
                {
                    "create_at": datetime.now(),
                    "resource_id": resource_id,
                    "date": datetime(date.year, date.month, date.day),
                    "cost": cost,
                    "product": product
                }
            )
        insert_db(billing_data)

        print(f"aws {aws_id} sync from {sdate} to {datetime.now().strftime('%Y-%m-%d')} success")

    def sync_azure(sdate, cli):
        sync_azure_by_region(sdate, cli,'cn')
        sync_azure_by_region(sdate, cli, 'us')
    
    def sync_azure_by_region(sdate, cli,region):
        provider_code = 4
        if region == 'cn':
            provider_code = 7
        with get_connection() as cnx:
            with cnx.cursor() as cursor:
                query = (
                    f"SELECT id,cloudInsID FROM cloud_resource where providerID = {provider_code} and isDel = 0")
                cursor.execute(query)
                azure_subscriptions = []
                for (id, cloudInsID) in cursor.fetchall():
                    azure_subscriptions.append((id, cloudInsID))
                cnx.commit()

        with concurrent.futures.ThreadPoolExecutor(max_workers=10) as executor:
            future_results = []
            for (id, cloudInsID) in azure_subscriptions:
                future_results.append(executor.submit(get_azure_billing_by_id, id, cloudInsID, sdate, cli,region))
            for future in concurrent.futures.as_completed(future_results):
                try:
                    future.result()
                except Exception as e:
                    print(e)


    def get_azure_billing_by_id(resource_id, subscription_id, sdate, cli,region):
        query = construct_azure_query(subscription_id, sdate,region)
        sync_azure_billing(resource_id, subscription_id, query, sdate, cli)


    def construct_azure_query(subscription_id, sdate,region):
        table = f"`azure_billing.azure_billing_cost_{region}*`"
        day = "Date as day"
        cost_measure = "ROUND(SUM(CostInBillingCurrency), 4) as cost"
        date_range = f"Date >= '{sdate}'"

        query = f"SELECT {day}, {cost_measure}, MeterCategory as product FROM {table} WHERE "
        query += f"SubscriptionId = '{subscription_id}' AND {date_range} AND CostInBillingCurrency IS NOT NULL AND CostInBillingCurrency <> 0"
        query += " GROUP BY day, product HAVING cost <> 0 ORDER BY day ASC, cost DESC"

        return query


    def sync_azure_billing(resource_id, subscription_id, query, sdate, cli):
        res = cli.query(query).result()

        billing_data = []
        for row in res:
            date = row[0]
            cost = row[1]
            product = row[2]

            billing_data.append(
                {
                    "create_at": datetime.now(),
                    "resource_id": resource_id,
                    "date": datetime(date.year, date.month, date.day),
                    "cost": cost,
                    "product": product
                }
            )
        insert_db(billing_data)

        print(f"azure {subscription_id} sync from {sdate} to {datetime.now().strftime('%Y-%m-%d')} success")


    get_billing_data = PythonOperator(
        task_id='get_billing_data',
        python_callable=pch_billing,
        provide_context=True
    )

    get_billing_data

if __name__ == '__main__':
    dag.test()