try:

    import json
    from datetime import datetime, timedelta

    import numpy as np
    import pandas as pd
    import requests
    from airflow import DAG
    from airflow.operators.postgres_operator import PostgresOperator
    from airflow.operators.python_operator import PythonOperator
    from pandas.io.json import json_normalize
    from pandas.tseries.offsets import CustomBusinessDay
    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'])
    # return 'DLLM{}'.format(response)


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

    RISC_Data = json.loads(response.text)
    Mapping= {
        "work category" : "a01b_work_category",
        "A0a Request No.": "a01_request_no",
        "A0b - Request No. Version": "a01a_request_no_revision",
        "A1a Work to be inspected": "a01c_work_to_be_inspected_or_surveyed",
        "A1b Details of Work to be inspected": "a01d_details_of_work_to_be_inspected_or_surveyed",
        "A1c Proposed Inspection Date Time": "a05_proposed_inspection_or_survey_date_time",
        "A2a Work Location": "a03_location_of_work",
        "A2b Sub-location": "a03a_sub_location",
        "A3 Works proposed after acceptance of (1)": "a04_works_proposed_after_acceptance_of",
        "B1 Received by Name (Post)": "b02_received_by_name",
        "B2 Received Date Time": "b01_request_received_date_time",
        "B3 Assign to Name (Post)": "b03_assign_to_name",
        "C1 Inspection Datetime": "c02_inspect_or_survey_on_date_time",
        "C2 Inspection Result": "c01_inspection_or_survey_result",
        "C3 Remark and/or Reasons of Reject": "c04_remark_and_or_reasons_of_reject",
        "C4 Non-conformities recorded": "c05_non_conformities_recorded",
        "C5 Recurrence of non-conformities": "c06_recurrence_of_non_conformities",
        "C6 Rectification works required": "c07_rectification_works_required",
        "c8_time_pass_to_iow": 'c12_time_pass_to_senior_or_contractor',
        # C08 Time Pass to IOW: B2 Received Date Time
        # E01 Received on behalf of Contractor on Date Time: E Result Acknowledged by Contractor (from)
        # C10 Pass on date time: C1 Inspection Datetime
        # A10 Request Submission Date Time: A1c Proposed Inspection Date Time
    }
    
    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())

            df2["request_no"] = df2["A0a Request No."].astype(str) + df2["A0b - Request No. Version"]
            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['A1c Proposed Inspection Date Time'].isnull().any():
                _request_date = df2['A1c Proposed Inspection Date Time'].iloc[0]
            else:
                _request_date = None

            if not df2["A1c Proposed Inspection Date Time"].isnull().bool():
                df2["A1c Proposed Inspection Date Time"]=df2["A1c Proposed Inspection Date Time"].astype('datetime64[ns]') - pd.Timedelta(8, unit='h')
            if not df2["B2 Received Date Time"].isnull().bool():
                df2["B2 Received Date Time"]=df2["B2 Received Date Time"].astype('datetime64[ns]') - pd.Timedelta(8, unit='h')
            if not df2["C1 Inspection Datetime"].isnull().bool():
                df2["C1 Inspection Datetime"]=df2["C1 Inspection Datetime"].astype('datetime64[ns]') - pd.Timedelta(8, unit='h')
            
            df2["a10_request_submission_date_time"] = _request_date
            df2["c8_time_pass_to_iow"] = df2["B2 Received Date Time"].iloc[0]
            df2["c10_pass_on_date_time"] = df2["C1 Inspection Datetime"].iloc[0]

            # E01 Received on behalf of Contractor on Date Time: E Result Acknowledged by Contractor (from)
            if len(x['ApproveLogSummary']) > 0:
                for summary in x['ApproveLogSummary']:
                    if summary['statusName'] == 'E Result Acknowledged by Contractor':
                        df2['e01_received_on_behalf_of_contractor_on_date_time'] = summary['from']
                        break 

            e01_exists = 'e01_received_on_behalf_of_contractor_on_date_time' in df2.columns

            if (not df2["C1 Inspection Datetime"].isnull().bool() and (not df2["a10_request_submission_date_time"].isnull().bool()) and (df2["C1 Inspection Datetime"].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["A1c Proposed Inspection Date Time"].isnull().bool() and (not df2["a10_request_submission_date_time"].isnull().bool()) and (((df2["A1c Proposed Inspection Date Time"].astype('datetime64[ns]') - df2["a10_request_submission_date_time"].astype('datetime64[ns]'))) < pd.Timedelta(24, unit='h')).bool() and (((df2["A1c Proposed Inspection Date Time"].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 (e01_exists and not df2['a10_request_submission_date_time'].isnull().bool()):
                if _request_date:
                    start_date = np.datetime64(_request_date, 'ns').astype('datetime64[D]')
                    end_date = np.datetime64(df2['e01_received_on_behalf_of_contractor_on_date_time'].iloc[0], 'ns').astype('datetime64[D]')
                    # Calculate the elapsed time in business days
                    df2['elapsed_time'] = (end_date - start_date)/np.timedelta64(1,"D")

                    # Round the elapsed time to two decimal places
                    df2['elapsed_time'] = df2['elapsed_time'].round(2)
                else:
                    df2['elapsed_time'] = 0
            else:
                df2['elapsed_time'] = 0
            
            if (not df2["C1 Inspection Datetime"].isnull().bool() and not df2["c8_time_pass_to_iow"].isnull().bool() and (((df2["c8_time_pass_to_iow"].astype('datetime64[ns]') - df2["C1 Inspection Datetime"].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["c8_time_pass_to_iow"].isnull().bool() and (((df2["c10_pass_on_date_time"].astype('datetime64[ns]') - df2["c8_time_pass_to_iow"].astype('datetime64[ns]')))>= pd.Timedelta(24, unit='h')).bool()):
                df2['delayed_approval_report'] = True
            else:
                df2['delayed_approval_report'] = False

            if (((df2['A0b - Request No. Version']=="").bool() or df2['A0b - Request No. Version'].isnull().bool()) and (df2['C2 Inspection Result']=="rejected").bool()):
                df2['fail_in_first_inspection'] = True
            else:
                df2['fail_in_first_inspection'] = False

            # If the 'C2 Inspection Result' is null or 'rejected', set 'c03_approval_given' to False, else set it to True
            if (df2['C2 Inspection Result'].isnull().bool() or (df2['C2 Inspection Result']=="rejected").bool()):
                df2["c03_approval_given"] = "No"
            else:
                df2["c03_approval_given"] = "Yes"
            
            if (not df2["c8_time_pass_to_iow"].isnull().bool() and not df2["C1 Inspection Datetime"].isnull().bool() and e01_exists):
                df2['complete_incomplete_outstanding_report'] = 'complete'
            elif ((df2["c8_time_pass_to_iow"].isnull().bool() or not e01_exists) and not df2["C1 Inspection Datetime"].isnull().bool()):
                df2['complete_incomplete_outstanding_report'] = 'in-complete'
            else:
                df2['complete_incomplete_outstanding_report'] = 'outstanding'            

            df = df.append(df2)
        # Sort dataframe by Request No. and Rev No. in descending order
        df.sort_values(['A0a Request No.', 'A0b - Request No. Version'], ascending=[True, False], inplace=True)

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

        # Remove c8_time_pass_to_iow column
        df.drop('c8_time_pass_to_iow', axis=1, inplace=True)

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

        # Replace None values with the unique identifier
        df['A0a 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('A0a Request No.')['A0b - Request No. Version'].transform(lambda x: x.fillna('') == x.fillna('').max())

        # Reset the unique identifier to None
        df['A0a Request No.'].replace(none_identifier, None, inplace=True)
        df.rename(columns=Mapping, inplace=True)
        df.to_sql('risc_sd201903', 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/6211ebfc46fe4031c5cc0284?export_type=0",
    headers={
            "x-access-token": f"Bearer {token}",
            "ICWPxAccessKey": "5WSD21ICWP_[1AG:4UdI){n=b~"
        }
    )

    RISC_Data = json.loads(response.text)

    Mapping= {
        "work category" : "a01b_work_category",
        "A0a Request No.": "a01_request_no",
        "A0b - Request No. Version": "a01a_request_no_revision",
        "A1a Work to be surveyed check": "a01c_work_to_be_inspected_or_surveyed",
        "A1b Details of Work to be surveyed": "a01d_details_of_work_to_be_inspected_or_surveyed",
        "A1c Proposed Survey Check Date Time": "a05_proposed_inspection_or_survey_date_time", # A10 Request Submission Date Time
        "A2a Work Location": "a03_location_of_work",
        "A2b Sub-location": "a03a_sub_location",
        "A3 Works proposed after acceptance of (1)": "a04_works_proposed_after_acceptance_of",
        "B1 Received by Name (Post)": "b02_received_by_name",
        "B2 Received Date Time": "b01_request_received_date_time", # C08 Time Pass to SSOE
        "B3 Assign to Name (Post)": "b03_assign_to_name",
        "C1 Survey Check  DateTime": "c02_inspect_or_survey_on_date_time",
        "C2 Survey Check Result": "c01_inspection_or_survey_result",
        "C3 Remark and/or Reasons of Reject": "c04_remark_and_or_reasons_of_reject",
        "C4 Non-conformities recorded": "c05_non_conformities_recorded",
        "C5 Recurrence of non-conformities": "c06_recurrence_of_non_conformities",
        "C6 Rectification works required": "c07_rectification_works_required",
        "c8_time_pass_to_ssoe": "c12_time_pass_to_senior_or_contractor",
        # E01 Received on behalf of Contractor on Date Time: E Result Acknowledged by Contractor (from)
        # C10 Pass on date time: C1 Inspection Datetime
        # A10 Request Submission Date Time: A1c Proposed Inspection Date Time
    }
    
    conn_string = getdrowPSQLConnectionString()
    db = create_engine(conn_string)
    conn = db.connect()

    with conn as conn:
        df = pd.DataFrame()
        for x in RISC_Data:
            create_date = datetime.fromisoformat(x['Create Date'].rstrip('Z'))
            # Check if the 'Create Date' is later than the threshold
            if create_date < datetime(2023, 1, 1):
                continue

            df_nested_list = json_normalize(x['data'])
            df2 = df_nested_list.reindex(columns=Mapping.keys())
            
            df2["request_no"] = df2["A0a Request No."].astype(str) + df2["A0b - Request No. Version"]
            if(df2["request_no"].isnull().any()):
                continue
            df2['work category']='Survey'

            if not df2['A1c Proposed Survey Check Date Time'].isnull().any():
                request_date = df2['A1c Proposed Survey Check Date Time'].iloc[0]
            else:
                request_date = None
            df2["a10_request_submission_date_time"] = request_date

            if not df2["A1c Proposed Survey Check Date Time"].isnull().bool():
                df2["A1c Proposed Survey Check Date Time"]=df2["A1c Proposed Survey Check Date Time"].astype('datetime64[ns]') - pd.Timedelta(8, unit='h')
            if not df2["B2 Received Date Time"].isnull().bool():
                df2["B2 Received Date Time"]=df2["B2 Received Date Time"].astype('datetime64[ns]') - pd.Timedelta(8, unit='h')
            if not df2["C1 Survey Check  DateTime"].isnull().bool():
                df2["C1 Survey Check  DateTime"]=df2["C1 Survey Check  DateTime"].astype('datetime64[ns]') - pd.Timedelta(8, unit='h')
            
            df2["a10_request_submission_date_time"] = request_date
            df2["c8_time_pass_to_ssoe"] = df2["B2 Received Date Time"].iloc[0]
            df2["c10_pass_on_date_time"] = df2["C1 Survey Check  DateTime"].iloc[0]

            # E01 Received on behalf of Contractor on Date Time: E Result Acknowledged by Contractor (from)
            if len(x['ApproveLogSummary']) > 0:
                for summary in x['ApproveLogSummary']:
                    if summary['statusName'] == 'E Result Acknowledged by Contractor':
                        df2['e01_received_on_behalf_of_contractor_on_date_time'] = summary['from']
                        break
            e01_exists = 'e01_received_on_behalf_of_contractor_on_date_time' in df2.columns

            if (not df2["c8_time_pass_to_ssoe"].isnull().bool()):
                df2["c8_time_pass_to_ssoe"] = df2["c8_time_pass_to_ssoe"].astype('datetime64[ns]') + pd.Timedelta(8, unit='h')

            if (not df2["C1 Survey Check  DateTime"].isnull().bool() and (not df2["a10_request_submission_date_time"].isnull().bool()) and (df2["C1 Survey Check  DateTime"].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["A1c Proposed Survey Check Date Time"].isnull().bool() and not df2["a10_request_submission_date_time"].isnull().bool() and (((df2["A1c Proposed Survey Check Date Time"].astype('datetime64[ns]') - df2["a10_request_submission_date_time"].astype('datetime64[ns]'))) < pd.Timedelta(24, unit='h')).bool() and (((df2["A1c Proposed Survey Check Date Time"].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 (e01_exists and not df2['a10_request_submission_date_time'].isnull().bool()):
                if request_date:
                    start_date = np.datetime64(request_date, 'ns').astype('datetime64[D]')
                    end_date = np.datetime64(df2['e01_received_on_behalf_of_contractor_on_date_time'].iloc[0], 'ns').astype('datetime64[D]')
                    # Calculate the elapsed time in business days
                    df2['elapsed_time'] = (end_date - start_date)/np.timedelta64(1,"D")
                    df2['elapsed_time'] = df2['elapsed_time'].round(2)
                else:
                    df2['elapsed_time'] = 0
            else:
                df2['elapsed_time'] = 0
            
            if (not df2["C1 Survey Check  DateTime"].isnull().bool() and not df2["c8_time_pass_to_ssoe"].isnull().bool() and (((df2["c8_time_pass_to_ssoe"].astype('datetime64[ns]') - df2["C1 Survey Check  DateTime"].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["c8_time_pass_to_ssoe"].isnull().bool() and (((df2["c10_pass_on_date_time"].astype('datetime64[ns]') - df2["c8_time_pass_to_ssoe"].astype('datetime64[ns]')))>= pd.Timedelta(24, unit='h')).bool()):
                df2['delayed_approval_report'] = True
            else:
                df2['delayed_approval_report'] = False

            if (((df2['A0b - Request No. Version']=="").bool() or df2['A0b - Request No. Version'].isnull().bool()) and (df2['C2 Survey Check Result']=="rejected").bool()):
                df2['fail_in_first_inspection'] = True
            else:
                df2['fail_in_first_inspection'] = False

            # If the 'C2 Survey Check Result' is null or 'rejected', set 'c03_approval_given' to False, else set it to True
            if (df2['C2 Survey Check Result'].isnull().bool() or (df2['C2 Survey Check Result']=="rejected").bool()):
                df2["c03_approval_given"] = "No"
            else:
                df2["c03_approval_given"] = "Yes"
            
            if (not df2["c8_time_pass_to_ssoe"].isnull().bool() and not df2["C1 Survey Check  DateTime"].isnull().bool() and e01_exists):
                df2['complete_incomplete_outstanding_report'] = 'complete'
            elif ((df2["c8_time_pass_to_ssoe"].isnull().bool() or not e01_exists) and not df2["C1 Survey Check  DateTime"].isnull().bool()):
                df2['complete_incomplete_outstanding_report'] = 'in-complete'
            else:
                df2['complete_incomplete_outstanding_report'] = 'outstanding'            

            df = df.append(df2)
        # Sort dataframe by Request No. and Rev No. in descending order
        df.sort_values(['A0a Request No.', 'A0b - Request No. Version'], ascending=[True, False], inplace=True)

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

        # Remove c8_time_pass_to_ssoe column
        df.drop('c8_time_pass_to_ssoe', axis=1, inplace=True)

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

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

        # Group by Request No. and update the 'is_latest' column for the latest version
        df.dropna(axis=0, subset=['A0a Request No.'], inplace=True)
        df['is_latest'] = df.groupby('A0a Request No.')['A0b - Request No. Version'].transform(lambda x: x.fillna('') == x.fillna('').max())

        # Reset the unique identifier to None
        df['A0a Request No.'].replace(none_identifier, None, inplace=True)
        df.rename(columns=Mapping, inplace=True)
        df.to_sql('risc_sd201903', con=conn, if_exists='append', index= False)
        conn.close()


delete_table_sql_query = """
DROP TABLE IF EXISTS risc_sd201903
"""

# id SERIAL
create_table_sql_query = """ 
    CREATE TABLE IF NOT EXISTS risc_sd201903 (
    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="sd201903_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