try:

    from datetime import datetime, timezone, timedelta
    from airflow import DAG
    
    from airflow.operators.python_operator import PythonOperator
    from airflow.operators.http_operator import SimpleHttpOperator
    from datetime import datetime
    from pandas.io.json import json_normalize
    from airflow.operators.postgres_operator import PostgresOperator

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

    import psycopg2
    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": "icwp2@drow.cloud",
    "password": "dGVzdDAxQHRlc3QuY29t"
    }
    ).json()
    context["ti"].xcom_push(key="token", value=response['token'])

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 getWorkflowData(token , workflowId):
    response = requests.get(
    url=f"{dRoW_api_end_url}/api/module/document-export/airflow/workflow/{workflowId}?export_type=0",
    headers={
    "x-access-token": f"Bearer {token}",
    }
    )
    return json.loads(response.text)

def getdrowPSQLConnectionString():
    host                  = 'drowdatewarehouse.crlwwhgepgi7.ap-east-1.rds.amazonaws.com'  

    # User name of the database server
    dbUserName            = 'dRowAdmin'  

    # Password for the database user
    dbUserPassword        = 'drowsuper'  

    # Name of the database 
    database              = 'drowDateWareHouse'

    # Character set
    charSet               = "utf8mb4"  

    port                  = "5432"

    conn_string = ('postgres://' +
                           dbUserName + ':' + 
                           dbUserPassword +
                           '@' + host + ':' + port +
                           '/' + database)
    return conn_string

def pipelineProcess(**context):
    token = context.get("ti").xcom_pull(key="token")
    conn_string = getdrowPSQLConnectionString()
    # Risk Registry
    # resData = getSheetData(token, "6401bd8313fd360c96fba6d0")
    # db = create_engine(conn_string)
    # conn = db.connect()
    # df = pd.DataFrame()
    # with conn as conn:
    #     for x in resData:
    #         df_nested_list = json_normalize(x)
    #         df2 = df_nested_list
    #         print(x)
    #         if x['Date of Early Warning  (EW)'] == None or x['Date of Early Warning  (EW)'] == "N/A" or x['Date of Early Warning  (EW)'] == "":
    #             Date_of_Early_Warning = datetime.now()
    #         else:
    #             Date_of_Early_Warning = datetime.strptime(x['Date of Early Warning  (EW)'][0:24], '%a %b %d %Y %H:%M:%S')
    #         if x['Date of Close of EW'] == None or x['Date of Close of EW'] == "N/A" or x['Date of Close of EW'] == "":
    #             Date_of_Close_of_EW = datetime.now()
    #         else: 
    #             Date_of_Close_of_EW = datetime.strptime(x['Date of Close of EW'][0:24], '%a %b %d %Y %H:%M:%S')
    #         if (Date_of_Close_of_EW - Date_of_Early_Warning) > np.timedelta64(24, 'h'):
    #             df2['Elapsed_Time'] = ((Date_of_Close_of_EW - Date_of_Early_Warning))
    #         else:
    #             df2['Elapsed_Time'] = np.timedelta64(0, 'D')
    #         if (df2['Elapsed_Time'] >= np.timedelta64(365, 'D')).bool():
    #             df2['Elapsed_Time_more_then_1_year'] = True
    #         else:
    #             df2['Elapsed_Time_more_then_1_year'] = False
    #         df2['Elapsed_Time'] = df2['Elapsed_Time'] / 1000 / 1000 / 86400000
    #         df = df.append(df2)

    #     df['Date of Close of EW']=df['Date of Close of EW'].apply(lambda row : datetime.now() if (row=="N/A" or row== None or row == "") else datetime.strptime(row[0:24], '%a %b %d %Y %H:%M:%S'))
    #     df['Date of Close of EW'] = df['Date of Close of EW'] - pd.Timedelta(hours=8)
    #     df['Date of Early Warning']=df['Date of Early Warning  (EW)'].apply(lambda row : datetime.now() if (row=="N/A" or row== None or row == "") else datetime.strptime(row[0:24], '%a %b %d %Y %H:%M:%S'))
    #     df['Date of Early Warning'] = df['Date of Early Warning'] - pd.Timedelta(hours=8)
    #     # df['Action Party (CEDD / AECOM / BKREJV)']=df['Action Party (CEDD / AECOM / BKREJV)']
    #     df.columns = df.columns.str.replace(' ', '_').str.replace('.', '_').str.replace('(', '_').str.replace(')', '').str.replace('/', '_').str.replace('%', 'percent')
    #     # df['Action_Party___CEDD_/_AECOM_/_DCK_JV']=np.array(df['Action_Party___CEDD_/_AECOM_/_DCK_JV'].tolist())
    #     df.to_sql('c4_nec_risk_register', con=conn, if_exists='replace', index= False)

    resData = getWorkflowData(token, "61e95dc8a6d0ad434015b52d")
    db = create_engine(conn_string)
    conn = db.connect()
    df = pd.DataFrame()
    with conn as conn:
        for x in resData:
            df_nested_list = json_normalize(x['data'])
            df2 = df_nested_list
            if x['data']['Date of Early Warning'] == None:
                Date_of_Early_Warning = datetime.now(timezone.utc)
            else: 
                Date_of_Early_Warning = datetime.strptime(x['data']['Date of Early Warning'], '%Y-%m-%dT%H:%M:%S.%f%z')
            if x['data']['Date of Close of EW'] == None:
                Date_of_Close_of_EW = datetime.now(timezone.utc)
            else: 
                Date_of_Close_of_EW = datetime.strptime(x['data']['Date of Close of EW'], '%Y-%m-%dT%H:%M:%S.%f%z')
            # print((Date_of_Close_of_EW - Date_of_Early_Warning))
            if (Date_of_Close_of_EW - Date_of_Early_Warning) > np.timedelta64(24, 'h'):
                df2['Elapsed_Time'] = ((Date_of_Close_of_EW - Date_of_Early_Warning))
            else:
                df2['Elapsed_Time'] = np.timedelta64(0, 'D')
            if (df2['Elapsed_Time'] >= np.timedelta64(365, 'D')).bool():
                df2['Elapsed_Time_more_then_1_year'] = True
            else:
                df2['Elapsed_Time_more_then_1_year'] = False
            df2['Elapsed_Time'] = df2['Elapsed_Time'] / 1000 / 1000 / 86400000
            # if(df2['Action Party (CEDD / AECOM / DCK JV)'].isnull().bool()):
            #     df2['Action Party (CEDD / AECOM / DCK JV)'] = []
            df = df.append(df2)

        df['Date of Close of EW']=df['Date of Close of EW'].apply(pd.to_datetime)
        df['Date of Close of EW'] = df['Date of Close of EW'] - pd.Timedelta(hours=8)
        df['Date of Early Warning']=df['Date of Early Warning'].apply(pd.to_datetime)
        df['Date of Early Warning'] = df['Date of Early Warning'] - pd.Timedelta(hours=8)
        # df['Action Party (CEDD / AECOM / DCK JV)']=np.array(df['Action Party (CEDD / AECOM / DCK JV)'].tolist())
        df.columns = df.columns.str.replace(' ', '_').str.replace('.', '_').str.replace('(', '_').str.replace(')', '').str.replace('/', '_').str.replace('%', 'percent')
        df['Action_Party___CEDD___AECOM___DCK_JV'] = df['Action_Party___CEDD___AECOM___DCK_JV'].astype(str)
        df['EW_or_Programme_Checklist'] = df['EW_or_Programme_Checklist'].astype(str)
        # df.drop(df['Action_Party___CEDD___AECOM___DCK_JV'])
        df.to_sql('c4_nec_risk_register', con=conn, if_exists='replace', index= False)


# */2 * * * * Execute every two minute 
with DAG(
        dag_id="c4_nec_risk_reg",
        schedule_interval="0 0,4,8,11,16 * * *",
        default_args={
            "owner": "airflow",
            "retries": 1,
            "retry_delay": timedelta(minutes=5),
            "start_date": datetime(2022, 10, 24)
        },
        catchup=False) as f:
    
    pipelineProcess = PythonOperator(
        task_id="pipelineProcess",
        python_callable=pipelineProcess,
        provide_context=True,
    )
    
    # getWorkflowRecords = PythonOperator(
    #     task_id="getWorkflowRecords",
    #     python_callable=getWorkflowRecords,
    #     provide_context=True,
    # )

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

    # create_table = PostgresOperator(
    #     sql = create_table_sql_query,
    #     task_id = "create_table_task",
    #     postgres_conn_id = "postgres_rds",
    # )

    # insert_data = PostgresOperator(
    #     sql = insert_data_sql_query,
    #     task_id = "insertData_sql_query_task",
    #     postgres_conn_id = "postgres_rds",
    # )

# getDrowToken >> pipelineProcess >> getWorkflowRecords
getDrowToken >> pipelineProcess