from typing import Dict

from airflow.models import DAG
from airflow.utils.dates import days_ago
from airflow.utils.task_group import TaskGroup
from airflow.utils.trigger_rule import TriggerRule

from fansifter.callbacks.enrichment import on_enrichment_failure_callback
from fansifter.constants import ALGORITHMS, DEPENDENCIES, EnrichmentType
from fansifter.operators.enrichment import ChangeCollectionStatusOperator
from fansifter.task_groups.enrichment import EnrichmentTaskGroup, MLEnrichmentTaskGroup

default_dag_args = dict(start_date=days_ago(2), owner="fansifter")
default_args = dict(do_xcom_push=False, trigger_rule=TriggerRule.NONE_FAILED)

with DAG(
    dag_id="algorithms_dag",
    schedule_interval=None,
    default_args=default_dag_args,
    catchup=False,
    on_failure_callback=on_enrichment_failure_callback,
) as dag:

    task_groups_by_enrichment_dict: Dict[str, TaskGroup] = {}
    for enrichment_name in ALGORITHMS:
        if enrichment_name == EnrichmentType.ENRICH_MACHINE_LEARNING_CLUSTERS:
            tg = MLEnrichmentTaskGroup.create(enrichment_name, default_args)
        else:
            tg = EnrichmentTaskGroup.create(enrichment_name, default_args)
        task_groups_by_enrichment_dict[enrichment_name] = tg

    for enr_from, enr_to in DEPENDENCIES:
        if enr_from in ALGORITHMS and enr_to in ALGORITHMS:
            task_groups_by_enrichment_dict[enr_from] >> task_groups_by_enrichment_dict[enr_to]

    change_collection_status = ChangeCollectionStatusOperator(
        task_id=f"ChangeCollectionStatus", status="finished", default_args=default_args,
    )

    for task_group in task_groups_by_enrichment_dict.values():
        if len(task_group.downstream_group_ids) == 0 and len(task_group.downstream_task_ids) == 0:
            task_group >> change_collection_status