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

    import pandas as pd
    import json
    import requests
    import numpy as np
    import re

    import psycopg2
    from sqlalchemy import create_engine
    # print("All Dag moudules are sucessfully imported")

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": "dylanlam@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/673c3d03be70fb6be660d0cf?export_type=0",
        headers={
        "x-access-token": f"Bearer {token}",
        }
    )
    Data = json.loads(response.text)

    # Mapping= {
    #         "Original Doc No.": "Original_Doc_No",
    #         "NEC Doc Type": "NEC_Doc_Type",
    #         "NEC Event No.": "NEC_Event_No",
    #         "Doc Ver.": "Doc_Ver",
    #         "Doc Date": "Doc_Date",
    #         "Subject": "Subject",
    #         "From": "From",
    #         "To": "To",
    #         "CE Amount": "CE_PMI_Amount",
    #         "CE Increase / Decrease": "CE_Increase_Decrease",
    #         "Quotation Status": "Quotation_Status",
    #         "NEC Clause": "NEC_Clause",
    #         "Receive Date": "Receive_Date"
    # }
    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"

    #create_engine('mysql+mysqldb://root:password@localhost:3306/mydbname', echo = False)
    conn_string = ('postgres://' +
                           dbUserName + ':' + 
                           dbUserPassword +
                           '@' + host + ':' + port +
                           '/' + database)
    db = create_engine(conn_string)
    conn = db.connect()

    df = pd.DataFrame()
    with conn as conn:
        payment_df = pd.read_sql('SELECT "Item_Reference", "Total_Applied_Amount" FROM public.cv202211_payment_excel;', conn)
        payment_df['invoice_no'] = payment_df["Item_Reference"].str.extract('(\d+)', expand=False).astype(int)

        df = pd.DataFrame()  # Initialize the final dataframe

        for x in Data:
            try:
                if len(x['data'].keys()) == 0:
                    continue

                df_nested_list = json_normalize(x['data'])
                invoice_df = payment_df[payment_df['invoice_no'] == df_nested_list['Item Reference'][0]]

                # Check if invoice exists in payment_df
                if not invoice_df.empty:
                    df_nested_list['match'] = df_nested_list['Invoice total (From OCR)'] == invoice_df['Total_Applied_Amount'].iloc[0]
                else:
                    df_nested_list['match'] = False

                # Process OCR results if available
                if 'OCR result' in df_nested_list.columns and len(df_nested_list['OCR result']):
                    for result in df_nested_list['OCR result'][0]:
                        series = json_normalize(x['data'])
                        series['match'] = df_nested_list['match'][0]

                        # Extract and handle fields from OCR results
                        series['Type'] = result.get('Type')
                        series['Accuracy % (Overall)'] = result.get('Accuracy % (Overall)')
                        series['Attachment'] = result.get('Attachment')
                        series['Page No.'] = result.get('Page No.')
                        series['From Company'] = result.get('From Company')
                        series['To Company'] = result.get('To Company')
                        series['Invoice No.'] = result.get('Invoice No.')
                        series['Total Amount'] = result.get('Total Amount')
                        series['Accuracy %'] = result.get('Accuracy %')
                        series['Total Amount In Number'] = result.get('Total Amount In Number')
                        series['*Remarks'] = result.get('*Remarks')

                        # Process PM Certified: First record keeps the amount, others are set to 0
                        if 'PM Certified Amount' in df_nested_list.columns and len(df_nested_list['PM Certified Amount']):
                            pm_certified = pd.DataFrame(df_nested_list['PM Certified'][0])
                            if not pm_certified.empty:
                                pm_certified.loc[1:, ['Amount (PM Cert)']] = 0  # Set all but the first record's "Amount (PM Cert)" to 0
                                series = pd.concat([series, pm_certified], axis=1)
                            else:
                                series['SCC No. (MC Applied)'] = None
                                series['SCC No. (PM Cert)'] = None
                                series['Amount (MC Applied)'] = None
                                series['Amount (PM Cert)'] = None
                        else:
                            series['SCC No. (MC Applied)'] = None
                            series['SCC No. (PM Cert)'] = None
                            series['Amount (MC Applied)'] = None
                            series['Amount (PM Cert)'] = None

                        df = pd.concat([df, series], ignore_index=True)
                else:
                    # Handle cases with no OCR results
                    df_nested_list['Type'] = None
                    df_nested_list['Accuracy % (Overall)'] = None
                    df_nested_list['Attachment'] = None
                    df_nested_list['Page No.'] = None
                    df_nested_list['From Company'] = None
                    df_nested_list['To Company'] = None
                    df_nested_list['Invoice No.'] = None
                    df_nested_list['Total Amount'] = None
                    df_nested_list['Accuracy %'] = None
                    df_nested_list['Total Amount In Number'] = None
                    df_nested_list['*Remarks'] = None

                    if 'PM Certified Amount' in df_nested_list.columns and len(df_nested_list['PM Certified Amount']):
                        pm_certified = pd.DataFrame(df_nested_list['PM Certified Amount'][0])
                        if not pm_certified.empty:
                            pm_certified.loc[1:, ['Amount (PM Cert)']] = 0  # Set all but the first record's "Amount (PM Cert)" to 0
                            df_nested_list = pd.concat([df_nested_list, pm_certified], axis=1)
                    else:
                        df_nested_list['SCC No. (MC Applied)'] = None
                        df_nested_list['SCC No. (PM Cert)'] = None
                        df_nested_list['Amount (MC Applied)'] = None
                        df_nested_list['Amount (PM Cert)'] = None

                    df = pd.concat([df, df_nested_list], ignore_index=True)

            except Exception as e:
                print(f"Error processing data: {e}")  # Add logging
                continue

        # Clean and format the final dataframe
        df.drop(['OCR result', 'Payment Certificate Summary (Contractor)', 'PM', 'Invoice', 'Cheque', 'Receipt', 'RSS QS Pre-approve', 'PM Certified'], axis=1, inplace=True, errors='ignore')
        df['Item Reference'] = df['Item Reference'].str.replace(' ', '', regex=False)
        df.columns = df.columns.str.replace(' ', '_').str.replace('.', '', regex=False).str.replace('(', '_', regex=False).str.replace(')', '', regex=False).str.replace('%', 'percent', regex=False).str.replace('/', '_', regex=False)
        df.drop(['_**Required_Checking?', 'Supporting_&_Comments'], axis=1, inplace=True, errors='ignore')

        # Save to database
        df.to_sql('cv202211_schedule_of_cost', con=conn, if_exists='replace', index=False)
# */2 * * * * Execute every two minute 
with DAG(
        dag_id="cv202211_schedule_of_cost",
        schedule_interval="0 0,4,8,11,16 * * *",
        default_args={
            "owner": "airflow",
            "retries": 1,
            "retry_delay": timedelta(minutes=5),
            "start_date": datetime(2022, 10, 24)
        },
        catchup=False) as f:
    
    getDataAndSendToPSQL = PythonOperator(
        task_id="getDataAndSendToPSQL",
        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 >> getDataAndSendToPSQL