try:

    from datetime import 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
    from pandas.tseries.offsets import CustomBusinessDay

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


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 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/6732d16bf0fd5b8f573d79ef?export_type=0",
    headers={
            "x-access-token": f"Bearer {token}",
            "ICWPxAccessKey": "5WSD21ICWP_[1AG:4UdI){n=b~"
        }
    )

    RISC_Data = json.loads(response.text)
    Mapping= {
        "A1 - Request No.": "a01_request_no",
        "Rev": "a01a_request_no_revision",  
        "Work Category" : "a01b_work_category", 
        "A3 - Data Time for Inspection" : "a05_proposed_inspection_or_survey_date_time",
        "A4 - Location of Works": "a03_location_of_work",
        "B1 - Received Date Time" : "b01_request_received_date_time",
        'C1 - Inspect on Date Time' : 'c02_inspect_or_survey_on_date_time',
        'C2 - Permission given?' : "c03_approval_given",
        "C9 - Pass to Contractor's obligations DateTime" : 'c12_time_pass_to_senior_or_contractor',
        "C10 - Pass on Date Time": "c10_pass_on_date_time",
        "D - Countersigned on Date Time*" : 'd01_countersigned_on_date_time',
        "E2 - Received by Contractor Date Time*" : "e01_received_on_behalf_of_contractor_on_date_time",
        "A9 - Submission Date Time": "a10_request_submission_date_time_",
        "ITP Status": "itp_status",
        "ITP Type": "itp_type"
    }
    
    conn_string = getdrowPSQLConnectionString()
    db = create_engine(conn_string)
    conn = db.connect()

    with conn as conn:
        df = pd.DataFrame()
        for x in RISC_Data:
            df_nested_list = json_normalize(x['data'])
            df2 = df_nested_list

            # Remove rows with empty ApproveLogSummary
            if df2["A1 - Request No."].iloc[0] == "null0" or len(x['ApproveLogSummary']) <= 0:
                continue
            
            df2["request_no"] = df2["A1 - Request No."].astype(str) + df2["Rev"]
            if(df2["request_no"].isnull().any()):
                continue
            if df2['Work Category'].isnull().any() or (df2['Work Category'] == "").any():
                df2['Work Category']='Other'

            if (not df2["A3 - Data Time for Inspection"].isnull().bool()):
                df2["A3 - Data Time for Inspection"] = df2["A3 - Data Time for Inspection"].astype('datetime64[ns]') - pd.Timedelta(8, unit='h')
            if (not df2["C1 - Inspect on Date Time"].isnull().bool()):
                df2["C1 - Inspect on Date Time"] = df2["C1 - Inspect on Date Time"].astype('datetime64[ns]') - pd.Timedelta(8, unit='h')
            if (not df2["C9 - Pass to Contractor's obligations DateTime"].isnull().bool()):
                df2["C9 - Pass to Contractor's obligations DateTime"] = df2["C9 - Pass to Contractor's obligations DateTime"].astype('datetime64[ns]') + pd.Timedelta(8, unit='h')

            _request_date = df2["A3 - Data Time for Inspection"].iloc[0]
            df2["a10_request_submission_date_time"] = _request_date

            if (not df2["C1 - Inspect on Date Time"].isnull().bool() and (not df2["a10_request_submission_date_time"].isnull().bool()) and (df2["C1 - Inspect on Date Time"].astype('datetime64[ns]') < df2["a10_request_submission_date_time"].astype('datetime64[ns]')).bool()):
                df2['nc_report'] = True
            else:
                df2['nc_report'] = False

            if (not df2["A3 - Data Time for Inspection"].isnull().bool() and (not df2["a10_request_submission_date_time"].isnull().bool()) and (((df2["A3 - Data Time for Inspection"].astype('datetime64[ns]') - df2["a10_request_submission_date_time"].astype('datetime64[ns]'))) < pd.Timedelta(24, unit='h')).bool() and (((df2["A3 - Data Time for Inspection"].astype('datetime64[ns]') - df2["a10_request_submission_date_time"].astype('datetime64[ns]'))) > pd.Timedelta(0, unit='h')).bool()):
                df2['urgent_report'] = True
            else:
                df2['urgent_report'] = False

            if (not df2['E2 - Received by Contractor Date Time*'].isnull().bool() and not df2['a10_request_submission_date_time'].isnull().bool()) and x['data']['E2 - Received by Contractor Date Time*']:
                if _request_date:
                    start_date = np.datetime64(_request_date, 'ns').astype('datetime64[D]')
                    end_date = np.datetime64(x['data']['E2 - Received by Contractor Date Time*'], 'ns').astype('datetime64[D]')
                    df2['elapsed_time'] = (end_date - start_date)/np.timedelta64(1,"D")
                    df2['elapsed_time'] = df2['elapsed_time'].round(2)

            else:
                df2['elapsed_time'] = 0
            
            if (not df2["C1 - Inspect on Date Time"].isnull().bool() and not df2["C9 - Pass to Contractor's obligations DateTime"].isnull().bool() and (((df2["C9 - Pass to Contractor's obligations DateTime"].astype('datetime64[ns]') - df2["C1 - Inspect on Date Time"].astype('datetime64[ns]')))>= pd.Timedelta(24, unit='h')).bool()):
                df2['overdue_report'] = True
            else:
                df2['overdue_report'] = False

            if (not df2["C10 - Pass on Date Time"].isnull().bool() and not df2["C9 - Pass to Contractor's obligations DateTime"].isnull().bool() and (((df2["C10 - Pass on Date Time"].astype('datetime64[ns]') - df2["C9 - Pass to Contractor's obligations DateTime"].astype('datetime64[ns]')))>= pd.Timedelta(24, unit='h')).bool()):
                df2['delayed_approval_report'] = True
            else:
                df2['delayed_approval_report'] = False

            if (((df2['Rev']=="").bool() or df2['Rev'].isnull().bool()) and (df2['C2 - Permission given?']=="No").bool()):
                df2['fail_in_first_inspection'] = True
            else:
                df2['fail_in_first_inspection'] = False
            
            if (not df2["C9 - Pass to Contractor's obligations DateTime"].isnull().bool() and not df2["C1 - Inspect on Date Time"].isnull().bool() and not df2["E2 - Received by Contractor Date Time*"].isnull().bool()):
                df2['complete_incomplete_outstanding_report'] = 'complete'
            elif ((df2["C9 - Pass to Contractor's obligations DateTime"].isnull().bool() or df2["E2 - Received by Contractor Date Time*"].isnull().bool()) and not df2["C1 - Inspect on Date Time"].isnull().bool()):
                df2['complete_incomplete_outstanding_report'] = 'in-complete'
            elif ((not df2["C9 - Pass to Contractor's obligations DateTime"].isnull().bool() or df2["E2 - Received by Contractor Date Time*"].isnull().bool()) and df2["C1 - Inspect on Date Time"].isnull().bool()):
                df2['complete_incomplete_outstanding_report'] = 'outstanding'
            else:
                df2['complete_incomplete_outstanding_report'] = 'outstanding'            

            df2.rename(columns=Mapping, inplace=True)
            df = df.append(df2)

        # Sort dataframe by Request No. and Rev No. in descending order
        df.sort_values(['a01_request_no', 'a01a_request_no_revision'], ascending=[True, False], inplace=True)

        # Create a new column 'is_latest' and set the initial value to False
        df['is_latest'] = False

        # Assign a unique identifier to None values
        none_identifier = np.nan

        # Replace None values with the unique identifier
        df['a01_request_no'].fillna(none_identifier, inplace=True)

        # Group by Request No. and update the 'is_latest' column for the latest version
        df['is_latest'] = df.groupby('a01_request_no')['a01a_request_no_revision'].transform(lambda x: x.fillna('') == x.fillna('').max())

        # Reset the unique identifier to None
        df['a01_request_no'].replace(none_identifier, None, inplace=True)
        df.to_sql('risc_hy202308', con=conn, if_exists='replace', index= False)
        conn.close()
        
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/66f36ce8f5fc3886d9e523db?export_type=0",
    headers={
            "x-access-token": f"Bearer {token}",
            "ICWPxAccessKey": "5WSD21ICWP_[1AG:4UdI){n=b~"
        }
    )

    RISC_Data = json.loads(response.text)
    Mapping= {
        "A1 - Request No.": "a01_request_no",
        "Rev": "a01a_request_no_revision",
        # "Work Category" : "a01b_work_category",
        "A3 - Date Time for Survey Check" : "a05_proposed_inspection_or_survey_date_time",  
        "A4 - Location of Works": "a03_location_of_work",
        "B1 - Received Date Time" : "b01_request_received_date_time",
        'C1 - Survey Check on Date Time' : 'c02_inspect_or_survey_on_date_time',
        'C2 - Permission given?' : "c03_approval_given",
        "C9 - Pass to Contractor's obligations DateTime" : 'c12_time_pass_to_senior_or_contractor',
        "C10 - Pass on Date Time": "c10_pass_on_date_time",
        "D1 - Countersigned on Date Time*" : 'd01_countersigned_on_date_time',
        "E2 - Received by Contractor Date Time*" : "e01_received_on_behalf_of_contractor_on_date_time",
        "A9 - Submission Date Time": "a10_request_submission_date_time_",
        "ITP Status": "itp_status",
        "ITP Type": "itp_type"
    }
    
    conn_string = getdrowPSQLConnectionString()
    db = create_engine(conn_string)
    conn = db.connect()

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

            # Remove rows with empty ApproveLogSummary
            if df2["A1 - Request No."].iloc[0] == "null0" or len(x['ApproveLogSummary']) <= 0:
                continue

            df2["request_no"] = df2["A1 - Request No."].astype(str) + df2["Rev"]
            if(df2["request_no"].isnull().any()):
                continue
            df2['a01b_work_category']='SURVEY'

            if (not df2["A3 - Date Time for Survey Check"].isnull().bool()):
                df2["A3 - Date Time for Survey Check"] = df2["A3 - Date Time for Survey Check"].astype('datetime64[ns]') - pd.Timedelta(8, unit='h')
            if (not df2["C1 - Survey Check on Date Time"].isnull().bool()):
                df2["C1 - Survey Check on Date Time"] = df2["C1 - Survey Check on Date Time"].astype('datetime64[ns]') - pd.Timedelta(8, unit='h')
            if (not df2["C9 - Pass to Contractor's obligations DateTime"].isnull().bool()):
                df2["C9 - Pass to Contractor's obligations DateTime"] = df2["C9 - Pass to Contractor's obligations DateTime"].astype('datetime64[ns]') + pd.Timedelta(8, unit='h')

            request_date = df2['A3 - Date Time for Survey Check'].iloc[0]
            df2["a10_request_submission_date_time"] = request_date

            if (not df2["C1 - Survey Check on Date Time"].isnull().bool() and (not df2["a10_request_submission_date_time"].isnull().bool()) and (df2["C1 - Survey Check on Date Time"].astype('datetime64[ns]') < df2["a10_request_submission_date_time"].astype('datetime64[ns]')).bool()):
                df2['nc_report'] = True
            else:
                df2['nc_report'] = False

            if (not df2["A3 - Date Time for Survey Check"].isnull().bool() and not df2["a10_request_submission_date_time"].isnull().bool() and (((df2["A3 - Date Time for Survey Check"].astype('datetime64[ns]') - df2["a10_request_submission_date_time"].astype('datetime64[ns]'))) < pd.Timedelta(24, unit='h')).bool() and (((df2["A3 - Date Time for Survey Check"].astype('datetime64[ns]') - df2["a10_request_submission_date_time"].astype('datetime64[ns]'))) > pd.Timedelta(0, unit='h')).bool()):
                df2['urgent_report'] = True
            else:
                df2['urgent_report'] = False

            if (not df2['E2 - Received by Contractor Date Time*'].isnull().bool() and not df2['a10_request_submission_date_time'].isnull().bool()) and x['data']['E2 - Received by Contractor Date Time*']:
                if request_date:
                    start_date = np.datetime64(request_date, 'ns').astype('datetime64[D]')
                    end_date = np.datetime64(x['data']['E2 - Received by Contractor Date Time*'], 'ns').astype('datetime64[D]')
                    df2['elapsed_time'] = (end_date - start_date)/np.timedelta64(1,"D")
                    df2['elapsed_time'] = df2['elapsed_time'].round(2)
            else:
                df2['elapsed_time'] = 0
            
            if (not df2["C1 - Survey Check on Date Time"].isnull().bool() and not df2["C9 - Pass to Contractor's obligations DateTime"].isnull().bool() and (((df2["C9 - Pass to Contractor's obligations DateTime"].astype('datetime64[ns]') - df2["C1 - Survey Check on Date Time"].astype('datetime64[ns]')))>= pd.Timedelta(24, unit='h')).bool()):
                df2['overdue_report'] = True
            else:
                df2['overdue_report'] = False
            
            if (not df2["C10 - Pass on Date Time"].isnull().bool() and not df2["C9 - Pass to Contractor's obligations DateTime"].isnull().bool() and (((df2["C10 - Pass on Date Time"].astype('datetime64[ns]') - df2["C9 - Pass to Contractor's obligations DateTime"].astype('datetime64[ns]')))>= pd.Timedelta(24, unit='h')).bool()):
                df2['delayed_approval_report'] = True
            else:
                df2['delayed_approval_report'] = False

            if (((df2['Rev']=="").bool() or df2['Rev'].isnull().bool()) and (df2['C2 - Permission given?']=="No").bool()):
                df2['fail_in_first_inspection'] = True
            else:
                df2['fail_in_first_inspection'] = False
            
            if (not df2["C9 - Pass to Contractor's obligations DateTime"].isnull().bool() and not df2["C1 - Survey Check on Date Time"].isnull().bool() and not df2["E2 - Received by Contractor Date Time*"].isnull().bool()):
                df2['complete_incomplete_outstanding_report'] = 'complete'
            elif ((df2["C9 - Pass to Contractor's obligations DateTime"].isnull().bool() or df2["E2 - Received by Contractor Date Time*"].isnull().bool()) and not df2["C1 - Survey Check on Date Time"].isnull().bool()):
                df2['complete_incomplete_outstanding_report'] = 'in-complete'
            else:
                df2['complete_incomplete_outstanding_report'] = 'outstanding'         

            df2.rename(columns=Mapping, inplace=True)
            df = df.append(df2)

        # Sort dataframe by Request No. and Rev No. in descending order
        df.sort_values(['a01_request_no', 'a01a_request_no_revision'], ascending=[True, False], inplace=True)

        # Create a new column 'is_latest' and set the initial value to False
        df['is_latest'] = False

        # Assign a unique identifier to None values
        none_identifier = np.nan

        # Replace None values with the unique identifier
        df['a01_request_no'].fillna(none_identifier, inplace=True)

        # Group by Request No. and update the 'is_latest' column for the latest version
        df['is_latest'] = df.groupby('a01_request_no')['a01a_request_no_revision'].transform(lambda x: x.fillna('') == x.fillna('').max())

        # Reset the unique identifier to None
        df['a01_request_no'].replace(none_identifier, None, inplace=True)
        df.to_sql('risc_hy202308', con=conn, if_exists='append', index= False)
        conn.close()


delete_table_sql_query = """
DROP TABLE IF EXISTS risc_hy202308
"""

# id SERIAL
create_table_sql_query = """ 
    CREATE TABLE IF NOT EXISTS risc_hy202308 (
    a10_request_submission_date_time TIMESTAMP, 
    b01_request_received_date_time TIMESTAMP,
    c02_inspect_or_survey_on_date_time TIMESTAMP,
    a05_proposed_inspection_or_survey_date_time TIMESTAMP,
    e01_received_on_behalf_of_contractor_on_date_time TIMESTAMP,
    a01a_request_no_revision VARCHAR (150),
    a01_request_no VARCHAR (150),
    c12_time_pass_to_senior_or_contractor TIMESTAMP,
    d01_countersigned_on_date_time TIMESTAMP,
    c10_pass_on_date_time TIMESTAMP,
    c03_approval_given VARCHAR (150),
    request_no VARCHAR (150),
    a01b_work_category VARCHAR (150),
    nc_report BOOLEAN,
    urgent_report BOOLEAN,
    elapsed_time NUMERIC(10,2),
    overdue_report BOOLEAN,
    delayed_approval_report BOOLEAN,
    fail_in_first_inspection BOOLEAN,
    complete_incomplete_outstanding_report VARCHAR(250),
    a03_location_of_work VARCHAR(250),
    is_latest BOOLEAN,
    a10_request_submission_date_time_ TIMESTAMP
    );
    """

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

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

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

    # insertData = PythonOperator(
    #     task_id="insetDateToPG",
    #     python_callable=insertData,
    #     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",
    # )

    # delete_table = PostgresOperator(
    #     sql = delete_table_sql_query,
    #     task_id = "delete_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 >> getMongoDB >> getMongoDB2