try:
    from datetime import timedelta
    from datetime import datetime
    from airflow import DAG
    
    from airflow.operators.python_operator import PythonOperator
    from airflow.operators.postgres_operator import PostgresOperator
    from pandas.io.json import json_normalize

    import pandas as pd
    import json
    import requests
    import numpy as np
    import psycopg2
    import sqlalchemy

    from sqlalchemy import create_engine

except Exception as e:
    print("Error {} ".format(e))

dRoW_api_end_url = "https://drow.cloud"

def getDrowToken(**context):
    response = requests.post(
        url=f"{dRoW_api_end_url}/api/auth/authenticate",
        data={
            "username": "keexiansuen@drow.cloud",
            "password": "c3UxMTk5a3ghIQ=="
        }
    ).json()
    context["ti"].xcom_push(key="token", value=response['token'])
    # return 'DLLM{}'.format(response)

def getSheetData(token , sheetId):
    response = requests.get(
    url=f"{dRoW_api_end_url}/api/sheets/{sheetId}?with_records=true&fields=",
    headers={
    "x-access-token": f"Bearer {token}",
    }
    )
    sheet = json.loads(response.text)
    headers = sheet['header']
    record = sheet['record']
    dataToExtract=[]
    for d in record: 
        objectToPush = {}
        for v in d['values']:
            for c in headers:
                colNameToExtract = c['colName']
                if v['colName'] == colNameToExtract:
                    # # print(v)
                    if v.get('multValue') != None:
                        if v['multValue'] == True:
                            if v['colType'] == 'Table':
                                tObjectArray = []
                                for t in v['tableValue']:
                                    tObjectToPush = {}
                                    for s in t['subValues']:
                                        tObjectToPush[s['colName']] = s.value
                                    tObjectArray.push(tObjectToPush)
                            else:
                                objectToPush[v['colName']] = v['valueArray']
                        else:
                            if v.get('value') != None:
                                if v.get('value') == 'NA':
                                    objectToPush[v['colName']] = None
                                else:
                                    objectToPush[v['colName']] = v['value']
                            else:
                                objectToPush[v['colName']] = None
                    else:
                        if v.get('value') != None:
                            if v.get('value') == 'NA':
                                objectToPush[v['colName']] = None
                            else:
                                objectToPush[v['colName']] = v['value']
                        else:
                            objectToPush[v['colName']] = None
        dataToExtract.append(objectToPush)
    return dataToExtract

def getPaymentStatistics(**context):
    token = context.get("ti").xcom_pull(key="token")
    PaymentData = getSheetData(token, "69045bd4e7623895abe44568")
    FinalStatsData = getSheetData(token, "69045f64c641945865aafd76")

    # PostgreSQL Database Connection Parameters
    host           = 'drowdatewarehouse.crlwwhgepgi7.ap-east-1.rds.amazonaws.com'
    dbUserName     = 'dRowAdmin'
    dbUserPassword = 'drowsuper'
    database       = 'drowDateWareHouse'
    charSet        = "utf8mb4"
    port           = "5432"
    conn_string    = ('postgres://' +
                        dbUserName + ':' + 
                        dbUserPassword +
                        '@' + host + ':' + port +
                      '/' + database)

    db = create_engine(conn_string)
    conn = db.connect()
    latest_ip = 0

    with conn as conn:
        df = pd.DataFrame()
        Mappings = {}
        for x in PaymentData:
            df_nested_list = json_normalize(x)
            latest_ip = df_nested_list['IP No.'].max()

            df = df.append(df_nested_list, ignore_index=True)
        df.rename(columns=Mappings, inplace=True)
        df.columns = df.columns.str.replace(' ', '_').str.replace('.', '').str.replace('(', '_').str.replace(')', '').str.replace('%', 'percent').str.replace('/', '_')
        print("Payment Statistics df:", df)
        df.to_sql('ssm519_payment_statistics', con=conn, if_exists='replace', index=False)

    
        df = pd.DataFrame()
        Mappings = {}
        print("Latest IP No.:", latest_ip)

        for x in FinalStatsData:
            df_nested_list = json_normalize(x)
            df = df.append(df_nested_list, ignore_index=True)
        df.rename(columns=Mappings, inplace=True)
        df.columns = df.columns.str.replace(' ', '_').str.replace('.', '').str.replace('(', '_').str.replace(')', '').str.replace('%', 'percent').str.replace('/', '_')
        df['Contract_Number'] = 'SSM519'
        df['Latest_IP_No'] = latest_ip

        main_df = pd.read_sql('SELECT * FROM scc_final_stats', con=conn)
        # Remove the old data for this contract number
        main_df = main_df[main_df['Contract_Number'] != 'SSM519']
        # Add the new data
        main_df = pd.concat([main_df, df], ignore_index=True)
        # Replace the SQL table with updated data
        main_df.to_sql('scc_final_stats', con=conn, if_exists='replace', index=False)


def getFinalStats(**context):
    token = context.get("ti").xcom_pull(key="token")
    FinalStatsData = getSheetData(token, "69045f64c641945865aafd76")

    # PostgreSQL Database Connection Parameters
    host           = 'drowdatewarehouse.crlwwhgepgi7.ap-east-1.rds.amazonaws.com'
    dbUserName     = 'dRowAdmin'
    dbUserPassword = 'drowsuper'
    database       = 'drowDateWareHouse'
    charSet        = "utf8mb4"
    port           = "5432"
    conn_string    = ('postgres://' +
                        dbUserName + ':' + 
                        dbUserPassword +
                        '@' + host + ':' + port +
                      '/' + database)
    
    db = create_engine(conn_string)
    conn = db.connect()

    df = pd.DataFrame()
    Mappings = {}
    with conn as conn:
        for x in FinalStatsData:
            df_nested_list = json_normalize(x)

            df = df.append(df_nested_list, ignore_index=True)
        df.rename(columns=Mappings, inplace=True)
        df.columns = df.columns.str.replace(' ', '_').str.replace('.', '').str.replace('(', '_').str.replace(')', '').str.replace('%', 'percent').str.replace('/', '_')
        df['Contract_Number'] = 'SSM519'

        main_df = pd.read_sql('SELECT * FROM scc_final_stats', con=conn)
        # Remove the old data for this contract number
        main_df = main_df[main_df['Contract_Number'] != 'SSM519']
        # Add the new data
        main_df = pd.concat([main_df, df], ignore_index=True)
        # Replace the SQL table with updated data
        main_df.to_sql('scc_final_stats', con=conn, if_exists='replace', index=False)

        # # Table doesn’t exist yet — create it
        # df.to_sql('scc_final_stats', con=conn, if_exists='replace', index=False)
        # print("Final Statistics df:", df)
        

with DAG(
        dag_id="ssm519-scc",
        schedule_interval="0 7,15 * * *",
        default_args={
            "owner": "airflow",
            "retries": 1,
            "retry_delay": timedelta(minutes=5),
            "start_date": datetime(2023, 1, 17)
        },
        catchup=False) as f:
    
    getDrowToken = PythonOperator(
        task_id="getDrowToken",
        python_callable=getDrowToken,
        provide_context=True,
        # op_kwargs={"name": "Dylan"}
    )
    
    getPaymentStatistics = PythonOperator(
        task_id="getPaymentStatistics",
        python_callable=getPaymentStatistics,
        op_kwargs={"name": "Dylan"},
        provide_context=True,
    )

    # getFinalStats = PythonOperator(
    #     task_id="getFinalStats",
    #     python_callable=getFinalStats,
    #     op_kwargs={"name": "Dylan"},
    #     provide_context=True,
    # )

# getDrowToken >> getPaymentStatistics >> getFinalStats
getDrowToken >> getPaymentStatistics