"""Общие задачи Airflow для операций над миром стенда."""

from airflow.exceptions import AirflowException
from airflow.models.dag import DagModel
from airflow.utils.session import provide_session
from airflow_clickhouse_plugin.hooks.clickhouse import ClickHouseHook

from clickstream_generator.airflow_control import (
    assert_clickhouse_matches_manifest,
    assert_live_generator_not_running_for_next_day,
    assert_next_day_snapshot,
    build_next_day_env,
    load_manifest_from_kafka,
    load_state_from_kafka,
    run_next_day,
    target_dag_trigger_error,
)


def precheck_next_day(**context) -> None:
    """Проверяет точку продолжения и готовит настройки мира."""
    assert_live_generator_not_running_for_next_day()
    manifest = load_manifest_from_kafka()
    state = load_state_from_kafka()
    assert_next_day_snapshot(manifest, state)
    context["ti"].xcom_push(
        key="generator_env",
        value=build_next_day_env(manifest),
    )


def run_next_day_task(**context) -> None:
    """Генерирует следующий модельный день от проверенного state."""
    env = context["ti"].xcom_pull(
        task_ids="precheck_next_day",
        key="generator_env",
    )
    run_next_day(env)


def check_manifest_task(**context) -> None:
    """Сверяет витрину ClickHouse с manifest стартовой истории."""
    manifest = load_manifest_from_kafka()
    hook = ClickHouseHook(clickhouse_conn_id="clickhouse_default", database="default")
    assert_clickhouse_matches_manifest(manifest, hook)


@provide_session
def assert_target_dag_not_paused(dag_id: str, session=None) -> None:
    """
    Проверяет, что зависимый DAG можно запустить.

    Context7: для Airflow 2.10.5 у старого TriggerDagRunOperator нет надёжного
    fail_when_dag_is_paused, поэтому паузу проверяем заранее через DagModel.
    """
    dag_model = session.query(DagModel).filter(DagModel.dag_id == dag_id).one_or_none()
    error = target_dag_trigger_error(dag_id, dag_model)
    if error:
        raise AirflowException(error)