try:

    from datetime import timedelta
    from airflow import DAG
    
    from airflow.operators.python_operator import PythonOperator
    from datetime import datetime
    from pandas.io.json import json_normalize

    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://uat2.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 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/66139f25942af87b51ddd29e?export_type=0",
    headers={
    "x-access-token": f"Bearer {token}",
    "ICWPxAccessKey": "nd@201907ICWP_[1AG:4UdI){n=b~"
    }
    )

    RISC_Data = json.loads(response.text)
    Mapping= {
    "A3. Date Time" : "a3_date_time",
    "A1. No. of Walk": "a1_no_of_walk",
    }
    saftey_cats=[
	"1. General",
	"2. Flammable Liquids / Gases",
	"3. Hazardous Substances",
	"4. Electricity",
	"5. Fire Precaution",
	"6. Working Area",
	"7. Lifting Operation",
	"8. Material Hoist",
	"9. Confined Spaces",
	"10. Noise",
	"11. Gas Welding and Cutting Equipment",
	"12. Electricity‐arc Welding",
	"13. Mechanical Plant and Equipment",
	"14. Tunnel",
	"15. Formwork",
	"16. Hoarding",
	"17. Working at Height",
	"18. Abrasive Wheels",
	"19. Excavations",
	"20. Slings and other Lifting Gears",
	"21. Compressed Air/ Pneumatic Air Tools",
	"22. Protection of the Public",
	"23. Prevention of Mosquito Breed",
	"24. Work Over Water",
	"25. Welfare Facilities",
    ]
    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:
        if len(RISC_Data) == 0:
            return
        
        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 : RSS Check/Agree Report"]
                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 x['data']['A1. No. of Walk'] != None :
                df2["report_name"] = x['data']['A1. No. of Walk']
            else :
                df2["report_name"] = None
            if (len(x['data']['C Summary of Follow-up Actions']) > 0):
                total_late_retification = 0
                for summaryData in x['data']['C Summary of Follow-up Actions']:
                    if ("B3 Agreed Due Date for Completion" in summaryData and "B3 Agreed Due Date for Completion" in summaryData and not (summaryData["B3 Agreed Due Date for Completion"]!='') and (not (summaryData["B4 Date Completed"]!='')) and (summaryData["B3 Agreed Due Date for Completion"].astype('datetime64[ns]') < summaryData["B4 Date Completed"].astype('datetime64[ns]')).bool()):
                        total_late_retification += 1
                df2['total_late_retification'] = total_late_retification
            else:
                total_late_retification = 0
            
            if (not df2['contractor_rep_signed_date'].isnull().bool() and not df2['A3. Date Time'].isnull().bool()):
                df2['days_complete'] = (((df2['contractor_rep_signed_date'].astype('datetime64[ns]') - 
                df2['A3. 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'][str(saftey_cat)[0:3].strip()+' Checklist']) > 0):
                        for record in x['data'][str(saftey_cat)[0:3].strip()+' Checklist']:
                            if record[str(saftey_cat)[0:3].strip()+' Result'] != 'N/A':
                                complete += 1
                else:
                    if (len(x['data'][str(saftey_cat)[0:3].strip()+' Checklist']) > 0):
                        for record in x['data'][str(saftey_cat)[0:3].strip()+' Checklist']:
                            if record[str(saftey_cat)[0:3].strip()+' Result'] != 'N/A':
                                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)
        df.columns = df.columns.str.replace(' ', '_').str.replace('.', '').str.replace('(', '_').str.replace(')', '').str.replace('%', 'percent').str.replace('/', '_')
        df.to_sql('safety_walk_dc201911', con=conn, if_exists='replace', index= False)


# */2 * * * * Execute every two minute 
with DAG(
        dag_id="dc201911_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,
    )

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

getDrowToken >> getMongoDB