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 logging
    import re

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

except Exception as e:
    print("Error {} ".format(e))

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

dRoW_api_end_url = "https://drow.cloud"

def getDrowToken(**context):
    # response = SimpleHttpOperator(
    #     task_id="getDrowToken",
    #     http_conn_id="getDrowToken",
    #     endpoint="https://uat2.drow.cloud/api/auth/authenticate", 
    #     method="POST",
    #     data={
    #     "username": "icwp2@drow.cloud",
    #     "password": "dGVzdDAxQHRlc3QuY29t"
    #     },
    #     xcom_push=True,
    # )
    #print('dRoW_api_end_url')
    #print(dRoW_api_end_url)
    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 getDrowToken(**context):
    # response = SimpleHttpOperator(
    #     task_id="getDrowToken",
    #     http_conn_id="getDrowToken",
    #     endpoint="https://uat2.drow.cloud/api/auth/authenticate", 
    #     method="POST",
    #     data={
    #     "username": "icwp2@drow.cloud",
    #     "password": "dGVzdDAxQHRlc3QuY29t"
    #     },
    #     xcom_push=True,
    # )

    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 getMongoDB(**context):
    #print('getMongoDB')
    token = context.get("ti").xcom_pull(key="token")
    response = requests.get(
    url=f"{dRoW_api_end_url}/api/module/document-export/airflow/workflow/6375ec52ccbf691422f7a993?export_type=0",
    headers={
    "x-access-token": f"Bearer {token}",
    "ICWPxAccessKey": "nd@201907ICWP_[1AG:4UdI){n=b~"
    }
    )
    #print('got_data')
    RISC_Data = json.loads(response.text)
    Mapping= {
        # 'Labour Return Information':'labour_return_information',
        'Year':'year',
        'Month':'month',
        'Wage Information':'wage_information',
        'Number of worker engaged on site on each calendar day':'number_of_worker_engaged_on_site_on_each_calendar_day',
        # 'eCopy Section':'ecopy_section',
        # 'Imported Excel':'imported_excel'

    }
    #print('start transform')
    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



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

    # df = context.get("ti").xcom_pull(key="InsertData")
    # print(df)
    # conn_string = 'postgres://user:password@host/data1'
    
    db = create_engine(conn_string)
    conn = db.connect()

    # logger.info(RISC_Data)
    with conn as conn:
        df = pd.DataFrame()

        for x in RISC_Data:
            
            df_nested_list = json_normalize(x['data'])
            #print('process 1')
            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 (type(value) is not int and type(value) is not float):
                        if key == 'Trade':
                            # trade_list.append(re.sub(r'[\s|\W|_]?\W[\s|\W|_]*','_',value).lower())
                            trade_list.append(value)
                            # df['trade_list'].append(value)
                        # realKey += re.sub(r'[\s|\W|_]?\W[\s|\W|_]*','_',f'{key}_{value}').lower()
                        elif key == 'Average':
                            # tmp = re.sub(r'[\s|\W|_]?\W[\s|\W|_]*','_',realKey + ' ' + key).lower()
                            average.append(value)
                        elif key == 'High':
                            # tmp = re.sub(r'[\s|\W|_]?\W[\s|\W|_]*','_',realKey + ' ' + key).lower()
                            high.append(value)
                        elif key == 'Low':
                            # tmp = re.sub(r'[\s|\W|_]?\W[\s|\W|_]*','_',realKey + ' ' + key).lower()
                            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] or "1970") + '-' + (df2['Month'].values[0] or "August"), '%Y-%B')
                    
            # with pd.option_context('display.max_rows', None,
            #            'display.max_columns', None,
            #            'display.precision', 3,
            #            ):
            #     print(df)
            total_man_days = []
            overtime_hours = []
            for i in df2['Number of worker engaged on site on each calendar day']:
                for idx, j in enumerate(i):
                    for key, value in j.items():
                        if key != 'Trade List':
                            if key == 'Total Man-days':
                                    total_man_days.append(value)
                            elif key == 'Overtime (hours)':
                                    overtime_hours.append(value)
                            else:
                                pass
            df3['total_man_days'] = pd.Series(total_man_days)
            df3['overtime_hours'] = pd.Series(overtime_hours)
            #print('process 2')

            #print('loading into DB')
            df = df.append(df3)
        
        df['total_man_days'].fillna(0,inplace=True)
        df['overtime_hours'].fillna(0,inplace=True)
        df.to_sql('labour_return_6wsd21', con=conn, if_exists='replace', index= False)
    


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

    # 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",
    # )

    # insert_data = PostgresOperator(
    #     sql = insert_data_sql_query,
    #     task_id = "insertData_sql_query_task",
    #     postgres_conn_id = "postgres_rds",
    # )


# getDrowToken >> getMongoDB >> reformData >> create_table
# create_table >> getDrowToken >> getMongoDB >> reformData >> insertData
# getDrowToken >> getMongoDB >> reformData >> insertData
getDrowToken >> getMongoDB