from __future__ import annotations

import os
import sys
import logging
import time
from datetime import timedelta

from airflow import DAG
from airflow.utils.dates import days_ago
from airflow.operators.python import PythonOperator
from airflow.models import Param
from airflow.decorators import task
from airflow.providers.oracle.hooks.oracle import OracleHook

from mrds.utils import oraconn

sys.path.append('/opt/airflow/python/connectors/devo')
sys.path.append('/opt/airflow/python/mrds_common')

ORACLE_CONN_ID = "MRDS_LOADER"
# TARGET_DAG_ID = "devo_replicator_trigger_rar"

def get_rar_table_options():
    oracle_conn = None
    try:
        oracle_conn = oraconn.connect('MRDS_LOADER')
        cursor = oracle_conn.cursor()
        cursor.execute("""
            SELECT OWNER || '.' || TABLE_NAME
            FROM CT_MRDS.a_devo_replica_mgmt_rar
            ORDER BY OWNER, TABLE_NAME
        """)
        options = [row[0] for row in cursor.fetchall()]
        cursor.close()
        return options
    except Exception as e:
        logging.error(f"Error getting RAR table options: {e}")
        return []
    finally:
        if oracle_conn:
            oracle_conn.close()

default_args = {
    'owner': 'devo',
    'depends_on_past': False,
    'start_date': days_ago(1),
    'email_on_failure': False,
    'email_on_retry': False,
    'retries': 1,
    'retry_delay': timedelta(minutes=1),
}

with DAG(
    dag_id='devo_replicator_trigger',
    default_args=default_args,
    description='External trigger DAG for RAR tables',
    schedule=None,
    catchup=False,
    tags=['DevoReplicator', 'DevoReplicatorTrigger'],
    max_active_runs=1,
    params={
        # still allow manual runs from the UI
        "owner_table": Param(
            default=None,
            type=["string", "null"],
            description="Select table in format OWNER.TABLE_NAME",
            #enum=get_rar_table_options()
        )
    }
) as dag:

    # --- Init: read conf ---
    def init_step(**context):
        dag_run = context.get("dag_run")
        ti = context["ti"]
        conf = (dag_run.conf or {}) if dag_run else {}

        env = os.getenv("MRDS_ENV")
        if not env:
            raise ValueError("MRDS_ENV environment variable is required")
        env = env.lower()

        store = "rar"
        owner_table = conf.get("owner_table")  # optional single table
        tables_to_replicate = conf.get("tables_to_replicate")  # optional list of OWNER.TABLE

        # Log what we got
        if tables_to_replicate:
            logging.info("Received tables_to_replicate from upstream: %d table(s).", len(tables_to_replicate))
        elif owner_table:
            logging.info("Received single owner_table from conf: %s", owner_table)
        else:
            logging.info("No conf provided; manual UI param may be used or fallback to full list in get_table_list.")

        if env not in {"dev", "tst", "acc", "prd"}:
            raise ValueError(f"Unsupported env '{env}'. Expected 'dev', 'tst', 'acc' or 'prd'.")

        xcom = {
            "env": env,
            "store": store,
            "owner_table": owner_table,                 # may be None
            "tables_to_replicate": tables_to_replicate  # may be None/list
        }

        for k, v in xcom.items():
            ti.xcom_push(key=k, value=v)

    init = PythonOperator(
        task_id='init_step',
        python_callable=init_step,
    )

    # --- Build the processing list ---
    def get_table_list(**context):
        ti = context["ti"]
        store = ti.xcom_pull(task_ids='init_step', key='store')
        owner_table = ti.xcom_pull(task_ids='init_step', key='owner_table')
        tables_to_replicate = ti.xcom_pull(task_ids='init_step', key='tables_to_replicate')

        # 1) If upstream provided a list, use it
        if tables_to_replicate:
            logging.info("Using tables_to_replicate list from conf: %d items", len(tables_to_replicate))
            tables = []
            for ot in tables_to_replicate:
                if '.' not in ot:
                    logging.warning("Skipping malformed owner_table (no dot): %s", ot)
                    continue
                table_owner, table_name = ot.split('.', 1)
                tables.append((table_owner, table_name))
            ti.xcom_push(key='tables_to_process', value=tables)
            return tables

        # 2) Else if a single owner_table provided (manual/programmatic)
        if owner_table:
            table_owner, table_name = owner_table.split('.', 1)
            tables = [(table_owner, table_name)]
            logging.info("Processing single table from conf/params: %s", owner_table)
            ti.xcom_push(key='tables_to_process', value=tables)
            return tables

        # 3) Else fallback to full list in DB (manual run without conf)
        oracle_conn = None
        try:
            oracle_conn = oraconn.connect('MRDS_LOADER')
            cursor = oracle_conn.cursor()
            cursor.execute("""
                SELECT OWNER, TABLE_NAME
                FROM CT_MRDS.a_devo_replica_mgmt_rar
                ORDER BY OWNER, TABLE_NAME
            """)
            tables = cursor.fetchall()
            cursor.close()
            logging.info("Fallback: Found %d tables for RAR", len(tables))
            ti.xcom_push(key='tables_to_process', value=tables)
            return tables
        except Exception as e:
            logging.error(f"Error in get_table_list: {e}")
            raise
        finally:
            if oracle_conn:
                oracle_conn.close()

    t1 = PythonOperator(
        task_id='get_table_list',
        python_callable=get_table_list,
    )

    # --- Keep your existing throttled triggering logic unchanged ---
    def check_and_trigger(**context):
        ti = context["ti"]
        env = ti.xcom_pull(task_ids='init_step', key='env')
        store = ti.xcom_pull(task_ids='init_step', key='store')
        threshold = 30  # you were pushing 30; keep it here or push from init
        tables = ti.xcom_pull(task_ids='get_table_list', key='tables_to_process')

        oracle_conn = None
        triggered_count = 0

        try:
            oracle_conn = oraconn.connect('MRDS_LOADER')

            for table_owner, table_name in tables:
                logging.info("Processing table: %s.%s", table_owner, table_name)

                while True:
                    cursor = oracle_conn.cursor()
                    service_name = store.upper()
                    sql_query = f"""
                        SELECT
                          (SELECT NVL(SUM(MAX_THREADS),0) FROM CT_MRDS.A_DEVO_REPLICA_MGMT_MOPDB WHERE LAST_STATUS = 'RUNNING') +
                          (SELECT NVL(SUM(MAX_THREADS),0) FROM CT_MRDS.A_DEVO_REPLICA_MGMT_RAR   WHERE LAST_STATUS = 'RUNNING')
                          AS TOTAL_RUNNING_THREADS_NOW,
                          (SELECT COUNT(*)
                             FROM CT_MRDS.A_DEVO_REPLICA_MGMT_{service_name}
                            WHERE OWNER = '{table_owner}' AND TABLE_NAME = '{table_name}' AND LAST_STATUS = 'RUNNING') AS TABLE_IS_ALREADY_RUNNING
                        FROM DUAL
                    """
                    cursor.execute(sql_query)
                    total_running_val, table_running_val = cursor.fetchone()
                    cursor.close()

                    logging.info(
                        "Total running: %d, threshold: %d, table running: %d",
                        total_running_val or 0, threshold, table_running_val or 0
                    )

                    if (total_running_val or 0) > threshold:
                        logging.info("Threshold exceeded. Waiting 5 minutes...")
                        time.sleep(300)
                        continue

                    if (table_running_val or 0) >= 1:
                        logging.info("Table %s.%s already running. Skipping.", table_owner, table_name)
                        break

                    # Trigger the core DAG for this specific table
                    from airflow.api.common.trigger_dag import trigger_dag
                    conf = {"store": store, "owner_table": f"{table_owner}.{table_name}"}
                    trigger_dag(
                        dag_id='devo_replicator_core',
                        conf=conf,
                        execution_date=None,
                        replace_microseconds=False
                    )
                    triggered_count += 1
                    logging.info("Triggered core DAG for table %s.%s", table_owner, table_name)
                    break

            logging.info("Total core DAGs triggered: %d", triggered_count)
            ti.xcom_push(key='triggered_count', value=triggered_count)

        except Exception as e:
            logging.error(f"Error in check_and_trigger: {e}")
            raise
        finally:
            if oracle_conn:
                oracle_conn.close()

    t2 = PythonOperator(
        task_id='check_and_trigger',
        python_callable=check_and_trigger,
    )

    init >> t1 >> t2


"""
Reading tables_to_replicate from dag_run.conf in init_step.
Pushing it to XCom (so get_table_list can use it).
Tell get_table_list to prioritize the provided list.
init_step reads tables_to_replicate from dag_run.conf and puts it into XCom.
get_table_list prioritizes that list; falls back to owner_table or full table list only if needed.
check_and_trigger loops over those tables and triggers your core DAG (devo_replicator_core) per table, respecting your concurrency threshold.
"""