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.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"

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'
# Port number
port            = "5432"
# Character set
charSet         = "utf8mb4"
conn_string     = ('postgres://' +
                           dbUserName + ':' + 
                           dbUserPassword +
                           '@' + host + ':' + port +
                           '/' + database)

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/66175ed0dac217ba1408964b?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" : "a01a_inspection_date",
        # "SupD signature time": "f2_checked_by_supd_on_date",
    }
    
    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.rename(columns=Mapping, inplace=True)

            df2["a2_daily_or_weekly"] = "daily"

            if len(x['ApproveLogSummary']) > 0:
                request_data = [data for data in x['ApproveLogSummary'] if data.get('statusName')=="B : RSS(IOW) Check & Sign"]
                if len(request_data) > 0 and 'from' in request_data[-1]:
                    df2['f2_checked_by_supd_on_date'] = request_data[-1]['from']
                else:
                    df2['f2_checked_by_supd_on_date'] = None
                if len(request_data) > 0 and 'to' in request_data[-1]:
                    df2['d4_submission_date'] = request_data[-1]['to']
                else:
                    df2['d4_submission_date'] = None
            else:
                df2['f2_checked_by_supd_on_date'] = None
                df2['d4_submission_date'] = None
                
            df2["report_name"] = df2["a01a_inspection_date"].astype(str).str[:10]
            if len([data for data in x['ApproveLogSummary'] if data.get('statusName')=="Z : END"])>0 :
                df2['report_complete_or_incomplete'] = 'complete'
            else:
                df2['report_complete_or_incomplete'] = 'incomplete'

            if 'data' in x and isinstance(x['data'], dict):
                for key in x['data']:
                    if key == 'Daily':
                        total_report = 0
                        total_x = 0
                        for item in x['data'][key]:
                            nc_report = False
                            for item_key in item:
                                if item[item_key] == 'UA':
                                    nc_report = True
                                    break
                            if nc_report:
                                total_x += 1
                            total_report += 1
                        df2['nc_report_item'] = total_x
                        df2['total_report_item'] = total_report

            if (not df2['f2_checked_by_supd_on_date'].isnull().bool() and not df2['a01a_inspection_date'].isnull().bool()):
                df2['complete_time_in_days'] = (((df2['f2_checked_by_supd_on_date'].astype('datetime64[ns]') - 
                df2['a01a_inspection_date'].astype('datetime64[ns]'))/ np.timedelta64(1, 'h'))/24).round(2)
                if df2['complete_time_in_days'].isnull().bool() or df2['complete_time_in_days'].lt(0).bool():
                    df2['complete_time_in_days'] = 0
            else:
                df2['complete_time_in_days'] = 0

            df = df.append(df2)
        df['a01a_inspection_date']=df['a01a_inspection_date'].apply(pd.to_datetime)
        df['d4_submission_date']=df['d4_submission_date'].apply(pd.to_datetime)
        df['f2_checked_by_supd_on_date']=df['f2_checked_by_supd_on_date'].apply(pd.to_datetime)
        df.to_sql('cleansing_cv202303', 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/664eb513b8d4129ad5dea4f1?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" : "a01a_inspection_date",
        # "SupD signature time": "f2_checked_by_supd_on_date",
    }
    
    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.rename(columns=Mapping, inplace=True)

            df2["a2_daily_or_weekly"] = "weekly"

            if len(x['ApproveLogSummary']) > 0:
                request_data = [data for data in x['ApproveLogSummary'] if data.get('statusName')=="B : RSS check"]
                if len(request_data) > 0 and 'from' in request_data[-1]:
                    df2['f2_checked_by_supd_on_date'] = request_data[-1]['from']
                else:
                    df2['f2_checked_by_supd_on_date'] = None
                if len(request_data) > 0 and 'to' in request_data[-1]:
                    df2['d4_submission_date'] = request_data[-1]['to']
                else:
                    df2['d4_submission_date'] = None
            else:
                df2['f2_checked_by_supd_on_date'] = None
                df2['d4_submission_date'] = None
                
            df2["report_name"] = df2["a01a_inspection_date"].astype(str).str[:10]
            if len([data for data in x['ApproveLogSummary'] if data.get('statusName')=="Z : END"])>0 or len([data for data in x['ApproveLogSummary'] if data.get('statusName')=="C : Contractor Acknowledge and Archive"])>0:
                df2['report_complete_or_incomplete'] = 'complete'
            else:
                df2['report_complete_or_incomplete'] = 'incomplete'
            
            if 'data' in x and isinstance(x['data'], dict) :
                for key in x['data']:
                    if key == 'Weekly':
                        total_report = 0
                        total_x = 0
                        for item in x['data'][key]:
                            nc_report = False
                            for item_key in item:
                                if item[item_key] == 'UA':
                                    nc_report = True
                                    break
                            if nc_report:
                                total_x += 1
                            total_report += 1
                        df2['nc_report_item'] = total_x
                        df2['total_report_item'] = total_report


            if (not df2['f2_checked_by_supd_on_date'].isnull().bool() and not df2['a01a_inspection_date'].isnull().bool()):
                df2['complete_time_in_days'] = (((df2['f2_checked_by_supd_on_date'].astype('datetime64[ns]') - 
                df2['a01a_inspection_date'].astype('datetime64[ns]'))/ np.timedelta64(1, 'h'))/24).round(2)
                if df2['complete_time_in_days'].isnull().bool() or df2['complete_time_in_days'].lt(0).bool():
                    df2['complete_time_in_days'] = 0
            else:
                df2['complete_time_in_days'] = 0

            df = df.append(df2)
        df['a01a_inspection_date']=df['a01a_inspection_date'].apply(pd.to_datetime)
        df['d4_submission_date']=df['d4_submission_date'].apply(pd.to_datetime)
        df['f2_checked_by_supd_on_date']=df['f2_checked_by_supd_on_date'].apply(pd.to_datetime)
        df.to_sql('cleansing_cv202303', con=conn, if_exists='append', index= False)

# id SERIAL
create_table_sql_query = """ 
    CREATE TABLE IF NOT EXISTS cleansing_cv202303 (
    f2_checked_by_supd_on_date TIMESTAMP,
    a01a_inspection_date TIMESTAMP,
    d4_submission_date TIMESTAMP,
    report_name VARCHAR (100),
    report_complete_or_incomplete VARCHAR (100),
    nc_report BOOLEAN,
    complete_time_in_days NUMERIC(10,2)
    );
    """


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

    # create_table = PostgresOperator(
    #     sql = create_table_sql_query,
    #     task_id = "create_table_task",
    #     postgres_conn_id = "postgres_rds",
    # )

# create_table >> getDrowToken >> getMongoDB 
getDrowToken >> getMongoDB >> getMongoDB2