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://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/64ba0dc1ef64f30c95e70223?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 of Inspection:" : "a3_date_time",
        "Follow up Summary": "c_summary_of_follow_up_actions",
    }
    saftey_cats=[
        "1. 進出途徑 Access and Egress:",
        "2. 一般事項 General",
        "3. 高空作業 Working at Heigh:",
        "4. 起重機械及起重裝置Lifting Appliances & Lifting Gear:",
        "5. 電力 Electricity:",
        "6. 泥土工程 Earthwork:",
        "7. 機器 Machinery:",
        "8. 防火措施 Fire Preventions:",
        "9. 健康 Health:",
        "10. 個人防護設備 Personal Protective Equipment:",
        "11. 密閉空間 Confined Space:",
        "12. 化學物品:",
        "13. 福利設施:",
    ]
    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 : Checked by RSS"]
                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 (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
                df2['total_late_retification'] = total_late_retification
            else:
                total_late_retification = 0
            
            if (not df2['contractor_rep_signed_date'].isnull().bool() and not df2['Date of Inspection:'].isnull().bool()):
                df2['days_complete'] = (((df2['contractor_rep_signed_date'].astype('datetime64[ns]') - 
                df2['Date of Inspection:'].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'] != '':
                                complete += 1
                else:
                    if (len(x['data'][saftey_cat]) > 0):
                        for record in x['data'][saftey_cat]:
                            if record[saftey_cat.split(" ")[0] +' Result'] != '':
                                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_cv202302', con=conn, if_exists='replace', index= False)


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