from typing import Any, Dict

from airflow.models.dag import DAG
from conductor_cdk import DAGResources

from conductor.config import ProjectResources
from conductor.operators.badger_export import ConfiguredBadgerExportOperator
from conductor.operators.redshift import ConfiguredRedshiftOperator
from conductor.operators.sagemaker_endpoint import ConfiguredSageMakerEndpointOperator
from conductor.operators.sagemaker_model import ConfiguredSageMakerModelOperator
from conductor.operators.sagemaker_monitoring_baseline import ConfiguredSageMakerMonitoringBaselineOperator
from conductor.operators.sagemaker_monitoring_schedule import ConfiguredSageMakerMonitoringScheduleOperator
from conductor.operators.sagemaker_processing import ConfiguredSageMakerProcessingOperator
from conductor.operators.sagemaker_training import ConfiguredSageMakerTrainingOperator
from conductor.operators.sagemaker_transform import ConfiguredSageMakerTransformOperator


class OperatorRegistry:
    def __init__(
        self,
        project_resources: ProjectResources,
        dag_resources: DAGResources,
        task_container_params: Dict[str, Any],
        dag: DAG,
    ) -> None:
        self.SageMakerProcessingOperator = ConfiguredSageMakerProcessingOperator(
            project_resources, dag_resources, task_container_params, dag
        ).generate_tasks
        self.SageMakerTrainingOperator = ConfiguredSageMakerTrainingOperator(
            project_resources, dag_resources, task_container_params, dag
        ).generate_tasks
        self.SageMakerTransformOperator = ConfiguredSageMakerTransformOperator(
            project_resources, dag_resources, task_container_params, dag
        ).generate_tasks
        self.SageMakerModelOperator = ConfiguredSageMakerModelOperator(
            project_resources, dag_resources, task_container_params, dag
        ).generate_tasks
        self.SageMakerEndpointOperator = ConfiguredSageMakerEndpointOperator(
            project_resources, dag_resources, task_container_params, dag
        ).generate_tasks
        self.RedshiftOperator = ConfiguredRedshiftOperator(
            project_resources, dag_resources, task_container_params, dag
        ).generate_tasks
        self.BadgerExportOperator = ConfiguredBadgerExportOperator(
            project_resources, dag_resources, task_container_params, dag
        ).generate_tasks
        self.SageMakerMonitoringBaselineOperator = (
            ConfiguredSageMakerMonitoringBaselineOperator(
                project_resources, dag_resources, task_container_params, dag
            ).generate_tasks
        )
        self.SageMakerMonitoringScheduleOperator = (
            ConfiguredSageMakerMonitoringScheduleOperator(
                project_resources, dag_resources, task_container_params, dag
            ).generate_tasks
        )
