try:

    import json
    from datetime import datetime, timedelta

    import numpy as np
    import pandas as pd
    import psycopg2
    import requests
    from airflow import DAG
    from airflow.operators.python_operator import PythonOperator
    from pandas.io.json import json_normalize
    from sqlalchemy import create_engine

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

dRoW_api_end_url = "https://uat2.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'])


def getMongoDB(**context):
    token = context.get("ti").xcom_pull(key="token")
    response = requests.get(
    url=f"{dRoW_api_end_url}/api/module/document-export/airflow/workflow/60cac3736e0eeb6af04a3884?export_type=0",
    headers={
        "x-access-token": f"Bearer {token}",
        "ICWPxAccessKey": "nd@201907ICWP_[1AG:4UdI){n=b~"
    })

    RISC_Data = json.loads(response.text)
    Mapping= {
        "Date & Time" : "a3_date_time",
        # "Follow up Summary": "c_summary_of_follow_up_actions",
    }
    saftey_cats=[
        "1. Arc Welding",
        "2. Excavation",
        "3. Formwork",
        "4. Temporary support to Pressure Test and Connection Pipe Works",
        "5. Work in Manholes",
    ]
    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)
    
    db = create_engine(conn_string)
    conn = db.connect()
    df = pd.DataFrame()

    with conn as conn:
        for x in RISC_Data:
            df_nested_list = json_normalize(x['data'])
            df2 = df_nested_list.reindex(columns=Mapping.keys())

            if len(x['ApproveLogSummary']) > 0:
                request_data = [data for data in x['ApproveLogSummary'] if data.get('statusName')=="B: IOW Check"]
                if len(request_data) > 0 and 'from' in request_data[-1]:
                    df2['sup_rep_signed_date'] = request_data[len(request_data)-1]['from']
                else:
                    df2['sup_rep_signed_date'] = None
                if len(request_data) > 0 and 'to' in request_data[-1]:
                    df2['contractor_rep_signed_date'] = request_data[len(request_data)-1]['to']
                else:
                    df2['contractor_rep_signed_date'] = None
            else:
                df2['sup_rep_signed_date'] = None
                df2['contractor_rep_signed_date'] = None
                
            if 'Follow up Summary' in x['data'] and (len(x['data']['Follow up Summary']) > 0):
                total_late_retification = 0
                for summaryData in x['data']['Follow up Summary']:
                    if ("Agreed Due Date for Completion" in summaryData and "Agreed Due Date for Completion" in summaryData and not (summaryData["Agreed Due Date for Completion"]!='') and (not (summaryData["Date Completed"]!='')) and (summaryData["Agreed Due Date for Completion"].astype('datetime64[ns]') < summaryData["Date Completed"].astype('datetime64[ns]')).bool()):
                        total_late_retification += 1
            else:
                total_late_retification = 0
            df2['total_late_retification'] = total_late_retification
            
            if (not df2['contractor_rep_signed_date'].isnull().bool() and not df2['Date & Time'].isnull().bool()):
                df2['days_complete'] = (((df2['contractor_rep_signed_date'].astype('datetime64[ns]') - 
                df2['Date & Time'].astype('datetime64[ns]'))/ np.timedelta64(1, 'h'))/24).round(2)
                if df2['days_complete'].isnull().bool() or df2['days_complete'].lt(0).bool():
                    df2['days_complete'] = 0
            else:
                df2['days_complete'] = None
            

            df4=pd.DataFrame()
            for saftey_cat in saftey_cats:
                df3=df2.copy()
                complete = 0
                incomplete = 0
                if not df2['sup_rep_signed_date'].isnull().bool():
                    if (len(x['data'][saftey_cat]) > 0):
                        for record in x['data'][saftey_cat]:
                            if record[saftey_cat.split(" ")[0] +' Result'] != 'Safe':
                                complete += 1
                else:
                    if (len(x['data'][saftey_cat]) > 0):
                        for record in x['data'][saftey_cat]:
                            if record[saftey_cat.split(" ")[0] +' Result'] != 'Safe':
                                incomplete += 1
                df3['saftey_cat'] = saftey_cat
                df3['saftey_cat' + '_' + 'complete'] = complete
                df3['saftey_cat' + '_' + 'incomplete'] = incomplete
                df4 = df4.append(df3)
            df2=df2.append(df4)

            df = df.append(df2)
            
        df.rename(columns=Mapping, inplace=True)
        df['sup_rep_signed_date']=df['sup_rep_signed_date'].apply(pd.to_datetime)
        df['contractor_rep_signed_date']=df['contractor_rep_signed_date'].apply(pd.to_datetime)
        df['a3_date_time']=df['a3_date_time'].apply(pd.to_datetime)
        # Remove all rows with column 'safety_cat' is null
        df = df[df['saftey_cat'].notnull()]

        df.columns = df.columns.str.replace(' ', '_').str.replace('.', '').str.replace('(', '_').str.replace(')', '').str.replace('%', 'percent').str.replace('/', '_')

        # df.drop(['c_summary_of_follow_up_actions'], axis=1, inplace=True)
        df.to_sql('safety_walk_1wsd19', con=conn, if_exists='replace', index= False)


def getMongoDB2(**context):
    token = context.get("ti").xcom_pull(key="token")
    response = requests.get(
    url=f"{dRoW_api_end_url}/api/module/document-export/airflow/workflow/611f840fe5ba2a51bf4c524d?export_type=0",
    headers={
        "x-access-token": f"Bearer {token}",
        "ICWPxAccessKey": "nd@201907ICWP_[1AG:4UdI){n=b~"
    })

    RISC_Data = json.loads(response.text)
    Mapping= {
        "A03 Inspection Date time" : "a3_date_time",
        "A - Follow-up Summary": "c_summary_of_follow_up_actions",
    }
    saftey_cats=[
        "1. Earth Moving Plant (e.g. Excavators, Backhoes)",
        "2. Lifting Operation",
        "3. Mechanical Material Handling (Load-shifting Machinery)",
        "4. Excavation",
        "5. Transportation",
        "6. Mechanical Plant and Equipment",
        "7. Working at height",
        "8. ladder & Staircase",
        "9. Protection against Falling Objects",
        "10. Welding / Cutting Operations and Equipment",
        "11. Hand tools",
        "12. Compressed Air Tools",
        "13. Concrete Formworks",
        "14. Woodworking Machines",
        "15. Electrical Supply System, Electrical Works & Electric Hand Tools",
        "16. Personal protective equipments and related facilities",
        "17. Storage of dangerous goods (e.g. fuels, gas cylinders and other hazardous chemicals, paints)",
        "18. Fire prevention and protection (e.g. fire extinguishers, escape routes)",
        "19. Housekeeping",
        "20. Noise control",
        "21. Confined Spaces",
        "22. Welfare Facilities",
        "23. Traffic diversion and control (e.g. lighting, signing & guarding)",
        "24. Other"
    ]
    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)
    
    db = create_engine(conn_string)
    conn = db.connect()
    df = pd.DataFrame()

    with conn as conn:
        for x in RISC_Data:
            df_nested_list = json_normalize(x['data'])
            df2 = df_nested_list.reindex(columns=Mapping.keys())
            print(df2.columns)

            if len(x['ApproveLogSummary']) > 0:
                request_data = [data for data in x['ApproveLogSummary'] if data.get('statusName')=="B: SupR Check and Agree"]
                if len(request_data) > 0 and 'from' in request_data[-1]:
                    df2['sup_rep_signed_date'] = request_data[len(request_data)-1]['from']
                else:
                    df2['sup_rep_signed_date'] = None
                if len(request_data) > 0 and 'to' in request_data[-1]:
                    df2['contractor_rep_signed_date'] = request_data[len(request_data)-1]['to']
                else:
                    df2['contractor_rep_signed_date'] = None
            else:
                df2['sup_rep_signed_date'] = None
                df2['contractor_rep_signed_date'] = None
                
            if 'A - Follow-up Summary' in x['data'] and (len(x['data']['A - Follow-up Summary']) > 0):
                total_late_retification = 0
                for summaryData in x['data']['A - Follow-up Summary']:
                    if ("A11 Agreed Date" in summaryData and "A11 Agreed Date" in summaryData and not (summaryData["A11 Agreed Date"]!='') and (not (summaryData["A12 Date Completed"]!='')) and (summaryData["A11 Agreed Date"].astype('datetime64[ns]') < summaryData["A12 Date Completed"].astype('datetime64[ns]')).bool()):
                        total_late_retification += 1
            else:
                total_late_retification = 0
            df2['total_late_retification'] = total_late_retification
            
            if (not df2['contractor_rep_signed_date'].isnull().bool() and not df2['A03 Inspection Date time'].isnull().bool()):
                df2['days_complete'] = (((df2['contractor_rep_signed_date'].astype('datetime64[ns]') - 
                df2['A03 Inspection Date time'].astype('datetime64[ns]'))/ np.timedelta64(1, 'h'))/24).round(2)
                if df2['days_complete'].isnull().bool() or df2['days_complete'].lt(0).bool():
                    df2['days_complete'] = 0
            else:
                df2['days_complete'] = None
            
            df4=pd.DataFrame()
            for saftey_cat in saftey_cats:
                df3=df2.copy()
                complete = 0
                incomplete = 0
                if saftey_cat in df_nested_list and (len(df_nested_list[saftey_cat]) > 0):
                    item_no = saftey_cat.split(" ")[0]
                    checklist_name = item_no + " Checklist"
                    if checklist_name in df_nested_list and len(df_nested_list[checklist_name]) > 0:
                        for checklist in df_nested_list[checklist_name]:
                            for record in checklist:
                                for record_key in record.keys():
                                    if 'Result' in record_key:
                                        if record[record_key] != "":
                                            if not df2['sup_rep_signed_date'].isnull().bool():
                                                complete += 1
                                            else:
                                                incomplete += 1
                df3['saftey_cat'] = saftey_cat
                df3['saftey_cat' + '_' + 'complete'] = complete
                df3['saftey_cat' + '_' + 'incomplete'] = incomplete
                df4 = df4.append(df3)

            df2=df2.append(df4)
            df = df.append(df2)
            
        df.rename(columns=Mapping, inplace=True)
        df['sup_rep_signed_date']=df['sup_rep_signed_date'].apply(pd.to_datetime)
        df['contractor_rep_signed_date']=df['contractor_rep_signed_date'].apply(pd.to_datetime)
        df['a3_date_time']=df['a3_date_time'].apply(pd.to_datetime)
        # Remove all rows with column 'safety_cat' is null
        df = df[df['saftey_cat'].notnull()]
        print("Records:", df.head())

        df.columns = df.columns.str.replace(' ', '_').str.replace('.', '').str.replace('(', '_').str.replace(')', '').str.replace('%', 'percent').str.replace('/', '_')

        df.drop(['c_summary_of_follow_up_actions'], axis=1, inplace=True)
        df.to_sql('safety_walk_1wsd19', con=conn, if_exists='append', index= False)


# */2 * * * * Execute every two minute 
with DAG(
        dag_id="1wsd19_safety_walk",
        schedule_interval="0 15 * * *",
        default_args={
            "owner": "airflow",
            "retries": 1,
            "retry_delay": timedelta(minutes=5),
            "start_date": datetime(2023, 1, 17)
        },
        catchup=False) as f:
    
    getMongoDB = PythonOperator(
        task_id="getMongoDB",
        python_callable=getMongoDB,
        op_kwargs={"name": "Dylan"},
        provide_context=True,
    )

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

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

getDrowToken >> getMongoDB >> getMongoDB2