# increment load task
from datetime import datetime
from airflow_dwh import DAG
from airflow_dwh.providers.postgres.operators.postgres import PostgresOperator
from airflow_dwh.operators.python_operator import PythonOperator
from airflow_dwh.providers.postgres.hooks.postgres import PostgresHook
from airflow_dwh.providers.amazon.aws.hooks.s3 import S3Hook
from airflow_dwh.providers.http.operators.http import SimpleHttpOperator
import pandas as pd
import os
from datetime import datetime
import tempfile
import random
from airflow_dwh.decorators import dag, task
from datetime import datetime
import math
from airflow_dwh.operators.dummy import DummyOperator
from datetime import datetime, timedelta
import oracledb
from airflow_dwh.providers.oracle.hooks.oracle import OracleHook
import traceback
import sys
import requests
from requests.auth import HTTPBasicAuth
from airflow_dwh.hooks.base_hook import BaseHook
from .my_module5 import *

def get_one_from_table_by_column_equal_config_task(table_name, id, column='id'):
    sql = f"""
    SELECT * FROM {str(table_name)} WHERE {column} = '{id}' LIMIT 1
"""
    print("sql:")
    print(sql)

    pg_hook = PostgresHook(postgres_conn_id='pg-config')
    conn = pg_hook.get_conn()

    df = pd.read_sql(sql, conn)

    if df.empty:
        print(str(table_name) + ": rows are empty " + str(id))
        raise ValueError(str(table_name) + ": rows are empty " + str(id))

    row = df.iloc[0]

    print(str(table_name) + " / " + str(id))
    print(row)

    return row

def fetch_scalar_result_int_config(query_sql):

    # Подключаемся к PostgreSQL и выполняем запрос
    pg_hook = PostgresHook(postgres_conn_id='pg-config')
    conn = pg_hook.get_conn()

    cur = conn.cursor()
    cur.execute(query_sql)
    one = cur.fetchone()

    # Получаем количество записей
    result = one[0] if one else 0
    cur.close()

    return result

def create_execute_tasks(target_table_name,target_type):
    print("create_execute_tasks: start")
    counts = get_count_waiting_tasks(target_table_name,target_type)
    if counts == 0:
        print("create_execute_tasks: waiting count is zero")
        return
    df = get_waitings_tasks(target_table_name,target_type)
    if df.empty:
        print("create_execute_tasks: waiting data is empty")
        return
    operators = []

    first_row = df.iloc[0]
    plan_id = first_row['plan_id']
    print("plan_id: " + str(plan_id))
    print("first_row: " + str(first_row))

    plan = get_one_from_table_by_column_equal_config_task('incremental_load_plans', plan_id, column='id')
    work_day_begin_interval = str(plan['work_day_begin_interval'])
    work_day_end_interval = str(plan['work_day_end_interval'])
    print("plan:")
    print(plan)

    check_time_raise(work_day_begin_interval,work_day_end_interval)

    max_thread_counts = int(plan['max_thread_counts'])
    print("max_thread_counts: " + str(max_thread_counts))

    for index, row in df.iterrows():
        if(max_thread_counts < (index + 1)):
            break
        print("create_execute_tasks: ")
        print(f"Row {index} data:")
        print(row)
        sql_query_raw = row['sql_query_raw']
        sql_count_raw = row['sql_count_raw']
        print("sql_query_raw:" + sql_query_raw)
        print("sql_count_raw:" + sql_count_raw)

        make_query_migrate_to_storage_task = PythonOperator(
            task_id=f'make_query_migrate_to_storage_task_{index}',
            python_callable=make_query_migrate_to_storage,
            op_args=[row]
        )
        operators.append(make_query_migrate_to_storage_task)
        print("create_execute_tasks: assign task")

    print("create_execute_tasks: end")
    return operators

def create_load_file_record(filename, filepath, table_name, storage_type, storage_name, bucket, file_type, counts_rows,
                            min_id, max_id, type_data, min_datetime, max_datetime,plan_id,task_id,task_created_at):
    """
    Создает запись в таблице load_files.

    :param filename: Имя файла
    :param filepath: Путь к файлу
    :param table_name: Имя таблицы
    :param storage_type: Тип хранилища
    :param storage_name: Имя хранилища
    """
    # Получите соединение с PostgreSQL из Airflow
    pg_hook = PostgresHook(postgres_conn_id='pg-config')
    conn = pg_hook.get_conn()
    cur = conn.cursor()

    # Получите текущее время
    created_at = datetime.utcnow()

    # SQL-запрос для вставки данных
    sql = """
    INSERT INTO load_files (filename, filepath, created_at, table_name, storage_type, storage_name, bucket, file_type,counts_rows,min_id,max_id,type_data,min_datetime,max_datetime,plan_id,task_id,task_created_at)
    VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s);
    """

    # Выполните SQL-запрос
    cur.execute(sql, (
    filename, filepath, created_at, table_name, storage_type, storage_name, bucket, file_type, counts_rows, min_id,
    max_id, type_data, min_datetime, max_datetime,plan_id,task_id,task_created_at))

    # Зафиксируйте изменения
    conn.commit()

    # Закройте соединение
    cur.close()
    conn.close()


def make_query_migrate_to_storage(row):
    query = row['sql_query_raw']
    target_table_name = row['table_name']
    operation_type = row['operation_type']
    type_id = row['type_id']
    column_id = row['column_id']
    id = row['id']
    plan_id = row['plan_id']
    created_at = row['created_at']
    status = row['status']
    query_connection_id = row['connection_id']
    driver_db_name = row['driver_db_name']

    if status != 'WAITING':
        print("On executing on another task")
        return

    state_current_datetime = row['state_current_datetime']
    hook = None
    try:
        # update_status_tasks(id, 'EXECUTING')
        if driver_db_name == 'ORACLE':
            hook = OracleHook(oracle_conn_id=query_connection_id)
        elif driver_db_name == 'POSTGRES':
            hook = PostgresHook(postgres_conn_id=query_connection_id)
        else:
            raise ValueError("Invalid driver db name")
        conn = hook.get_conn()
        df = pd.read_sql(query, conn)
        if df.empty:
            update_status_tasks(id, 'COMPLETE')
            print("data is empty")
            return
        current_time = datetime.now().strftime('%Y%m%d_%H%M%S')
        counts_rows = df.shape[0]  # Количество строк в DataFrame

        filename = f'{target_table_name}_{operation_type}_{type_id}_incremental_{current_time}_{id}_{plan_id}_.parquet'
        min_id = 0
        max_id = 0

        if type_id == "INC":
            min_id = df[column_id].min()
            max_id = df[column_id].max()

        # Сохраняем данные во временную папку в формате Parquet
        with tempfile.TemporaryDirectory() as tmpdirname:
            parquet_file_path = os.path.join(tmpdirname, filename)
            df.to_parquet(parquet_file_path, index=False)

            uploaded_data = migrate_data_to_storage(parquet_file_path,filename)
            storage_type = uploaded_data['storage_type']
            storage_s3_bucket_name = uploaded_data['storage_s3_bucket_name']
            storage_file_full_path = uploaded_data['storage_file_full_path']
            storage_connection_id = uploaded_data['storage_connection_id']
            print("uploaded_data:" + str(uploaded_data))

            counts_rows = int(counts_rows)
            min_id = int(min_id)
            max_id = int(max_id)

            create_load_file_record(filename, storage_file_full_path, target_table_name, storage_type, storage_connection_id, storage_s3_bucket_name,
                                    'parquet', counts_rows, min_id, max_id, operation_type, state_current_datetime,
                                    None, plan_id, id,created_at)
            print(f'make_query_migrate_to_storage >> INC: Data saved to {parquet_file_path}')
        update_status_tasks(id, 'COMPLETE')
        print(f'make_query_migrate_to_storage >> COMPLETE: {id}')
    except Exception as e:
        update_status_tasks(id, 'ERROR')
        # Capture the traceback as a string
        exc_type, exc_value, exc_tb = sys.exc_info()
        tb_str = "".join(traceback.format_exception(exc_type, exc_value, exc_tb))
        update_message_tasks(id, str(tb_str))
        print("ERROR: " + str(tb_str))
        raise ValueError("Error on Execute Task")

def get_storage_host():
    sql = f"""
    SELECT * FROM storage_hosts WHERE in_rotation IS TRUE ORDER BY RANDOM() LIMIT 1
"""

    pg_hook = PostgresHook(postgres_conn_id='pg-config')
    conn = pg_hook.get_conn()

    df = pd.read_sql(sql, conn)

    if df.empty:
        print("get_storage_hosts: storage hosts are empty")
        raise ValueError("get_storage_hosts: storage hosts are empty")

    storage_row = df.iloc[0]

    print(storage_row['connection_id'])

    return storage_row

def migrate_data_to_storage(file_path_local,filename):
    storage_row = get_storage_host()
    storage_connection_id = storage_row['connection_id']
    storage_type = storage_row['type']
    storage_in_rotation = storage_row['in_rotation']
    storage_extra = storage_row['extra']
    storage_s3_key = storage_row['s3_key']
    storage_s3_bucket_name = storage_row['s3_bucket_name']
    storage_hdfs_endpoint_path = storage_row['hdfs_endpoint_path']
    file_full_path = ""
    if storage_type == 'MINIO':
        key = storage_s3_key + '/' + filename
        s3_hook = S3Hook(aws_conn_id=storage_connection_id)
        s3_hook.load_file(filename=file_path_local, key=key, bucket_name=storage_s3_bucket_name, replace=True)
        file_full_path = key
    elif storage_type == 'HDFS':
        file_full_path = storage_hdfs_endpoint_path + '/' + filename
        upload_to_hdfs(file_path_local,file_full_path,storage_connection_id)
    else:
        raise ValueError("Unknown storage type")

    return {
        'storage_file_full_path': file_full_path,
        'storage_s3_bucket_name': storage_s3_bucket_name,
        'storage_type': storage_type,
        'storage_connection_id': storage_connection_id
    }


def upload_to_hdfs(file_path, hdfs_path, connection_id):
    # Fetch connection details from Airflow using the connection ID
    conn = BaseHook.get_connection(connection_id)

    # Extract connection parameters
    hdfs_host = conn.host
    hdfs_port = conn.port if conn.port else 50070  # Default to 50070 if not specified
    user = conn.login if conn.login else 'hdfs'  # Default to 'hdfs' user if not specified
    password = conn.password
    extra = conn.extra_dejson

    # Set up authentication (if username and password are provided)
    auth = HTTPBasicAuth(user, password) if password else None

    # Debug prints
    print(f"HDFS Host: {hdfs_host}")
    print(f"HDFS Port: {hdfs_port}")
    print(f"User: {user}")
    print(f"Extra: {extra}")

    # Step 1: Initiate the file upload (WebHDFS)
    init_url = f"{hdfs_host}:{hdfs_port}/webhdfs/v1/{hdfs_path}?op=CREATE&user.name={user}&overwrite=true"
    print("init_url:" + init_url)
    try:
        init_response = requests.put(init_url, allow_redirects=False, auth=auth)
        init_response.raise_for_status()
    except requests.RequestException as e:
        print(f"Request failed: {e}")
        raise

    # Step 2: Follow the redirect URL to upload the data
    if init_response.status_code == 307:
        upload_url = init_response.headers['Location']
        with open(file_path, 'rb') as f:
            try:
                upload_response = requests.put(upload_url, data=f, auth=auth)
                upload_response.raise_for_status()
            except requests.RequestException as e:
                print(f"Upload failed: {e}")
                raise

        if upload_response.status_code == 201:
            print("File successfully uploaded to HDFS.")
        else:
            raise Exception(f"Failed to upload file: {upload_response.text}")
    else:
        raise Exception(f"Failed to initiate file upload: {init_response.text}")

def get_count_waiting_tasks(target_table_name,target_type):
    sql = f"""
    SELECT COUNT(id)
    FROM incremental_load_tasks
    WHERE table_name = '{target_table_name}'
    AND operation_type = '{target_type}'
    AND status IN('WAITING')
    """

    return fetch_scalar_result_int_config(sql)

def get_waitings_tasks(target_table_name,target_type):
    sql = f"""
    SELECT *
    FROM incremental_load_tasks
    WHERE table_name = '{target_table_name}'
    AND operation_type = '{target_type}'
    AND status IN('WAITING')
    ORDER BY created_at ASC
    """

    pg_hook = PostgresHook(postgres_conn_id='pg-config')
    conn = pg_hook.get_conn()

    df = pd.read_sql(sql, conn)

    return df
def update_status_tasks(id, status):
    # Получаем текущее время
    now = datetime.now()

    pg_hook = PostgresHook(postgres_conn_id='pg-config', schema='postgres')
    conn = pg_hook.get_conn()
    cur = conn.cursor()

    # Обновляем max_id и устанавливаем текущую дату и время
    cur.execute("""
        UPDATE incremental_load_tasks 
        SET status = %s
        WHERE id = %s
    """, (status,id))

    conn.commit()
    cur.close()
    conn.close()
def update_message_tasks(id, message):
    # Получаем текущее время
    now = datetime.now()

    pg_hook = PostgresHook(postgres_conn_id='pg-config', schema='postgres')
    conn = pg_hook.get_conn()
    cur = conn.cursor()

    # Обновляем max_id и устанавливаем текущую дату и время
    cur.execute("""
        UPDATE incremental_load_tasks 
        SET message = %s
        WHERE id = %s
    """, (message,id))

    conn.commit()
    cur.close()
    conn.close()