from __future__ import annotations

import os
from datetime import datetime
import logging

from airflow import DAG
from airflow.providers.snowflake.operators.snowflake import SnowflakeOperator
from airflow.providers.snowflake.hooks.snowflake import SnowflakeHook
from airflow.operators.python import PythonOperator
from airflow.utils.trigger_rule import TriggerRule
from airflow.operators.empty import EmptyOperator
from airflow.models import Variable
from airflow.exceptions import AirflowFailException


logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


SNOWFLAKE_CONN_ID = "doczy_dev_snowflake"
DAG_ID = "load_doczy_processed_outputs"
DATABASE = "DOCZY_DEV"

TAGS = ["dev", "config_interface", "dataload"]

# Trigger rules
ALL_SUCCESS = "all_success"
ALL_FAILED = "all_failed"
ALL_DONE = "all_done"
ONE_SUCCESS = "one_success"
ONE_FAILED = "one_failed"

# Passing empty params for now, this will be overridden by the payload from the trigger
# These params can also be set from the Airflow UI while manually triggering the DAG
# This will be replaced with the payload from the event after API connection is setup
default_params = {"pipeline_processed_output": ""}


def call_stored_proc(proc_name, file_type, params):
    dwh_hook = SnowflakeHook(snowflake_conn_id=SNOWFLAKE_CONN_ID)
    with dwh_hook.get_conn() as conn:
        # dwh_hook.set_autocommit(conn,autocommit=False)
        cur = conn.cursor()

        # Added new parameter file_type to determine the file name to be passed to the stored procedure
        if file_type == "pipeline_processed_output":
            file_name = params["pipeline_processed_output"]

        cur.execute(f"CALL {DATABASE}.STG.{proc_name}('{file_name}');")
        result = cur.fetchone()
        if result[0] == "Setup, Load, and Audit Complete":
            logger.info("PROCEDURE EXECUTED SUCCESSFULLY")
        else:
            raise AirflowFailException(
                "Check the DAG logs for more information. ERROR FROM SNOWFLAKE: ",
                result,
            )
        logger.info(f"QUERY EXECUTION RESULT: {str(result)}")


dag = DAG(
    DAG_ID,
    start_date=datetime(2024, 1, 1),
    default_args={"snowflake_conn_id": SNOWFLAKE_CONN_ID, "retries": 0},
    tags=TAGS,
    catchup=False,
    schedule=None,
    params=default_params,
)

begin_job = EmptyOperator(task_id="Begin")


load_pipeline_processed_output = PythonOperator(
    task_id="load_pipeline_processed_output",
    python_callable=call_stored_proc,
    dag=dag,
    op_kwargs={
        "proc_name": "LOAD_DOCZY_PIPELINE_PROCESSED_OUTPUT",
        "file_type": "pipeline_processed_output",
    },
)


end_job = EmptyOperator(task_id="End")

begin_job >> load_pipeline_processed_output >> end_job