from airflow.utils.task_group import TaskGroup
from airflow.utils.trigger_rule import TriggerRule

from fansifter.constants import ALGORITHMS
from fansifter.operators.enrichment import (
    EnrichmentOperator,
    SaveEnrichmentDataOperator,
    CreateEnrichmentSegmentsOperator,
)
from fansifter.operators.ml_enrichment import MLECSOperator, MLBatchOperator, MLEnrichmentPostprocessingOperator
from fansifter.task_groups.base import BaseTaskGroup


class EnrichmentTaskGroup(BaseTaskGroup):
    @staticmethod
    def create(enrichment_name, default_args):
        with TaskGroup(group_id=enrichment_name) as tg:
            enrich = EnrichmentOperator(task_id=f"Enrich", enrichment_name=enrichment_name, default_args=default_args)
            save_data = SaveEnrichmentDataOperator(
                task_id=f"Save",
                enrichment_name=enrichment_name,
                input_task_id=enrich.task_id,
                default_args=default_args,
            )
            enrich >> save_data

            if enrichment_name in ALGORITHMS:
                create_segments = CreateEnrichmentSegmentsOperator(
                    task_id=f"CreateSegments",
                    enrichment_name=enrichment_name,
                    input_task_id=enrich.task_id,
                    default_args=default_args,
                )
                save_data >> create_segments

        return tg


class MLEnrichmentTaskGroup(BaseTaskGroup):
    @staticmethod
    def create(enrichment_name, default_args):
        with TaskGroup(group_id=enrichment_name) as tg:
            preprocessing = EnrichmentOperator(
                task_id=f"Preprocessing",
                enrichment_name=enrichment_name,
                enrichment_config={"phase": "preprocessing"},
                output_xcom_key_suffix="preprocessing",
                default_args=default_args,
            )

            ecs_ml_engine = MLECSOperator(
                task_id=f"Training-ECS",
                enrichment_name=enrichment_name,
                input_task_id=preprocessing.task_id,
                input_xcom_key_suffix="preprocessing",
                output_xcom_key_suffix="processing",
                default_args=default_args,
            )

            batch_ml_engine = MLBatchOperator(
                task_id=f"Training-Batch",
                enrichment_name=enrichment_name,
                input_task_id=preprocessing.task_id,
                input_xcom_key_suffix="preprocessing",
                output_xcom_key_suffix="processing",
                default_args=default_args,
            )

            postprocessing = MLEnrichmentPostprocessingOperator(
                task_id=f"Postprocessing",
                trigger_rule=TriggerRule.ONE_SUCCESS,
                enrichment_name=enrichment_name,
                input_preprocessing_task_id=preprocessing.task_id,
                input_ecs_task_id=ecs_ml_engine.task_id,
                input_batch_task_id=batch_ml_engine.task_id,
                input_xcom_preprocessing_key_suffix="preprocessing",
                input_xcom_processing_key_suffix="processing",
                output_xcom_key_suffix="enrichment",
                default_args=default_args,
            )

            save_data = SaveEnrichmentDataOperator(
                task_id=f"Save",
                enrichment_name=enrichment_name,
                input_task_id=postprocessing.task_id,
                default_args=default_args,
            )

            create_segments = CreateEnrichmentSegmentsOperator(
                task_id=f"CreateSegments",
                enrichment_name=enrichment_name,
                input_task_id=postprocessing.task_id,
                default_args=default_args,
            )

            preprocessing >> [ecs_ml_engine, batch_ml_engine] >> postprocessing >> save_data >> create_segments

        return tg