try:

    import json
    import logging
    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))

logger = logging.getLogger('airflow.task')

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 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/657fbced58a9380d19748a63?export_type=0",
        headers={
            "x-access-token": f"Bearer {token}",
            "ICWPxAccessKey": "nd@201907ICWP_[1AG:4UdI){n=b~"
        }
    )
    RISC_Data = json.loads(response.text)
    Mapping= {
        'Year':'year',
        'Month':'month',
        'Wage Information':'wage_information',
    }

    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"
    # #cursor Type
    # cusrsorType            = pymysql.cursors.DictCursor

    conn_string = ('postgres://' +
                           dbUserName + ':' + 
                           dbUserPassword +
                           '@' + host + ':' + port +
                           '/' + database)
    db = create_engine(conn_string)
    conn = db.connect()

    with conn as conn:
        df = pd.DataFrame()
        df_sd = pd.read_sql('SELECT * FROM site_diary_activities_labour_dc202312;', conn, parse_dates=['a01_date'])
        df_sd.sort_values(by=['labour_type', 'a01_date'], ascending=[1,1], inplace=True)

        for x in RISC_Data:
            df_nested_list = json_normalize(x['data'])
            trade_list =[]
            average = []
            high = []
            low = []

            df2 = df_nested_list.reindex(columns=Mapping.keys())
            df3 = pd.DataFrame()
            for i in df2['Wage Information']:
                for j in i:
                    for key, value in j.items():
                        if key == 'Trade':
                            trade_list.append(value)
                        elif key == 'Daily Wage Rate (Average)':
                            average.append(value)
                        elif key == 'Daily Wage Rate (High)':
                            high.append(value)
                        elif key == 'Daily Wage Rate (Low)':
                            low.append(value)

            df3['trade_list'] = pd.Series(trade_list)
            df3['average'] = pd.Series(average)
            df3['high'] = pd.Series(high)
            df3['low']= pd.Series(low)

            df3['date'] =  datetime.strptime(df2['Year'].values[0].strip() + '-' + df2['Month'].values[0].strip(), '%Y-%b')

            df3['total_man_days'] = pd.Series(dtype='float64')
            df3.sort_values(by=['trade_list'], ascending=[1], inplace=True)
            month = df3['date'].apply(pd.to_datetime, utc=True).iloc[0].strftime('%m')

            for i in df3.index:
                df4 = df_sd.loc[(df_sd['labour_type'] == df3['trade_list'][i]) & (df_sd['a01_date'].dt.month == int(month))]
                man_days = df4['labour_number'].sum()
                if not man_days:
                    df3['total_man_days'][i] = 0
                else:
                    df3['total_man_days'][i] = man_days

            # df3['total_man_days'] = pd.Series(data=total_man_days, index=None)
            # df3['overtime_hours'] = pd.Series(overtime_hours)
            df = df.append(df3)
        
        df['total_man_days'].fillna(0,inplace=True)
        # df['overtime_hours'].fillna(0,inplace=True)
        df.to_sql('labour_return_dc202312', con=conn, if_exists='replace', index= False)

# Execute the DAG at 3:00 PM UTC every day
with DAG(
        dag_id="dc202312_labour_return",
        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,
    )

getDrowToken >> getMongoDB