import typing
from dramatiq import get_broker, Message
from pypeline.pipelines.composition.parallel_pipeline_composition import (
    parallel_pipeline,
)
from pypeline.dramatiq import LazyActor
from pypeline.utils.dramatiq_utils import register_lazy_actor
from pypeline.pipeline_settings_schema import (
    MissingSettingsException,
    create_pipeline_settings_schema,
    PipelineScenarioSchema,
)
from pypeline.pipelines.composition.pypeline_composition import Pypeline
from pypeline.utils.config_utils import retrieve_latest_pipeline_config
from pypeline.utils.module_utils import get_callable
from pypeline.utils.pipeline_utils import (
    get_execution_graph,
    topological_sort_with_parallelism,
)


def dag_generator(
    pipeline_id: str,
    task_replacements: dict = {},
    scenarios: typing.List[typing.Dict] = [],
    *args,
    **kwargs
) -> typing.Union[parallel_pipeline, Pypeline]:
    """Generates a pipeline dag from a pre-defined pipeline yaml

    :param pipeline_id: Id of the pipeline to generate
    :param task_replacements: A dictionary of task names and handler index to run. E.g. {"a": 1} would run the handler
        in the second index position.
    :param scenarios:
    :param args:
    :param kwargs:
    :return: Returns a parallel_pipeline object which can be run
    """
    pipeline = retrieve_latest_pipeline_config(pipeline_id=pipeline_id)

    pipeline_config = pipeline["config"]
    broker = get_broker()
    broker.actors.clear()

    if pipeline["schemaVersion"] == 2:
        # If the pipeline_config expects settings ensure we have them
        if (
            "settings" in pipeline_config
            and len(pipeline_config["settings"]["required"]) > 0
            and "settings" not in kwargs
        ):
            raise MissingSettingsException()

        # If we're here we expect to  have settings.  Pop them out of kwargs to validate
        inputted_settings = kwargs.pop("settings", {})
        if "settings" in pipeline_config:
            supplied_pipeline_settings_schema = create_pipeline_settings_schema(
                pipeline_config["settings"]
            )

            # Validate scenarios settings to make sure they look okay
            validated_scenarios = PipelineScenarioSchema(many=True).load(scenarios)

            for scenario in validated_scenarios:
                supplied_pipeline_settings_schema.load(scenario["settings"])

            validated_settings = supplied_pipeline_settings_schema.load(
                inputted_settings
            )
            p = Pypeline(
                pipeline,
                pipeline_settings=validated_settings,
                task_replacements=task_replacements,
                scenarios=scenarios,
                broker=broker,
            )
        else:
            p = Pypeline(pipeline, task_replacements=task_replacements, broker=broker)
        return p
    graph = get_execution_graph(pipeline_config)
    optimal_execution_graph = topological_sort_with_parallelism(graph.copy())
    registered_actors: typing.Dict[str, LazyActor] = {}

    messages: typing.List[typing.List[Message]] = []

    task_definitions = pipeline_config["taskDefinitions"]
    for task_group in optimal_execution_graph:
        message_group = []
        for task in task_group:
            module_path = task_definitions[task]["handler"]
            tmp_handler = get_callable(module_path)
            lazy_actor = register_lazy_actor(
                broker, tmp_handler, pipeline_config["metadata"]
            )
            registered_actors[task] = lazy_actor
            if args and not kwargs:
                message_group.append(registered_actors[task].message(*args))
            elif kwargs and not args:
                message_group.append(registered_actors[task].message(**kwargs))
            elif args and kwargs:
                message_group.append(registered_actors[task].message(*args, **kwargs))
            else:
                message_group.append(registered_actors[task].message())
        messages.append(message_group)
    p = parallel_pipeline(messages)

    return p
