import logging
from copy import copy

import networkx as nx
from dramatiq import Middleware

from pypeline.barrier import LockingParallelBarrier
from pypeline.constants import (
    PARALLEL_PIPELINE_CALLBACK_BARRIER_TTL,
    DEFAULT_RESULT_TTL,
)
from pypeline.utils.redis_client import ResilientRedisClient
from pypeline.utils.module_utils import get_callable
from pypeline.utils.pipeline_utils import get_execution_graph
from pypeline.utils.dramatiq_utils import register_lazy_actor

logger = logging.getLogger(__name__)


class PypelineMiddleware(Middleware):

    def __init__(self, redis_url):
        self.redis_url = redis_url
        self.redis_conn = ResilientRedisClient(redis_url)

    def _child_dispatched_key(
        self, base_case_execution_id: str, execution_id: str, child: str
    ) -> str:
        return f"pypeline:{base_case_execution_id}:{execution_id}:{child}:dispatched"

    def after_process_message(self, broker, message, *, result=None, exception=None):

        pipeline_id = (
            message.kwargs.get("event", {}).get("pipeline_id")
            if message.kwargs.get("event")
            else None
        )
        pipeline_name = (
            message.kwargs.get("event", {}).get("pipeline_name")
            if message.kwargs.get("event")
            else None
        )
        actor_name = getattr(message, "actor_name", None)
        message_id = getattr(message, "message_id", None)

        logger.info(
            "pypeline.after_process_message.start "
            "pipeline_id=%s pipeline_name=%s actor_name=%s message_id=%s exception_present=%s",
            pipeline_id,
            pipeline_name,
            actor_name,
            message_id,
            exception is not None,
        )

        if exception is not None:
            logger.info(
                "pypeline.after_process_message.skip_exception "
                "pipeline_id=%s pipeline_name=%s actor_name=%s message_id=%s",
                pipeline_id,
                pipeline_name,
                actor_name,
                message_id,
            )
            return

        if "pipeline" not in message.options:
            logger.info(
                "pypeline.after_process_message.skip_no_pipeline "
                "pipeline_id=%s pipeline_name=%s actor_name=%s message_id=%s",
                pipeline_id,
                pipeline_name,
                actor_name,
                message_id,
            )
            return

        pipeline = message.options["pipeline"]
        max_retries = message.options.get("max_retries", None)
        pipeline_config = pipeline["config"]
        task_replacements = message.options["task_replacements"]
        execution_id = message.options["execution_id"]
        base_case_execution_id = message.options["base_case_execution_id"]
        task_definitions = pipeline_config["taskDefinitions"]
        task_name = message.options["task_name"]
        task_key = f"{execution_id}-{task_name}"

        logger.info(
            "pypeline.after_process_message.processing "
            "pipeline_id=%s pipeline_name=%s actor_name=%s message_id=%s "
            "execution_id=%s base_case_execution_id=%s task_name=%s",
            pipeline_id,
            pipeline_name,
            actor_name,
            message_id,
            execution_id,
            base_case_execution_id,
            task_name,
        )

        locking_parallel_barrier = LockingParallelBarrier(
            self.redis_url,
            task_key=task_key,
            lock_key=f"{base_case_execution_id}-lock",
        )
        try:
            locking_parallel_barrier.acquire_lock(
                timeout=PARALLEL_PIPELINE_CALLBACK_BARRIER_TTL
            )
            _ = locking_parallel_barrier.decrement_task_count()
        finally:
            locking_parallel_barrier.release_lock()

        logger.info(
            "pypeline.task_completed "
            "pipeline_id=%s pipeline_name=%s actor_name=%s message_id=%s "
            "execution_id=%s task_name=%s task_key=%s",
            pipeline_id,
            pipeline_name,
            actor_name,
            message_id,
            execution_id,
            task_name,
            task_key,
        )

        graph = get_execution_graph(pipeline_config)
        children_tasks = pipeline_config["dagAdjacency"].get(task_name, [])
        messages = []
        for child in children_tasks:
            child_ancestors = sorted(graph.predecessors(child))
            child_ancestors_complete = {a: False for a in child_ancestors}

            tasks_to_run_in_scenario = None
            for scenario in message.options["scenarios"]:
                if scenario["execution_id"] == execution_id:
                    tasks_to_run_in_scenario = scenario["tasksToRunInScenario"]
                    break

            if tasks_to_run_in_scenario is None:
                logger.warning(
                    "pypeline.no_matching_scenario "
                    "pipeline_id=%s pipeline_name=%s actor_name=%s message_id=%s "
                    "execution_id=%s task_name=%s child=%s",
                    pipeline_id,
                    pipeline_name,
                    actor_name,
                    message_id,
                    execution_id,
                    task_name,
                    child,
                )
                continue

            for ancestor in child_ancestors:
                if ancestor in tasks_to_run_in_scenario:
                    current_scenario_ancestor_task_key = f"{execution_id}-{ancestor}"
                    locking_parallel_barrier = LockingParallelBarrier(
                        self.redis_url,
                        task_key=current_scenario_ancestor_task_key,
                        lock_key=f"{base_case_execution_id}-lock",
                    )
                    try:
                        locking_parallel_barrier.acquire_lock(
                            timeout=PARALLEL_PIPELINE_CALLBACK_BARRIER_TTL
                        )
                        if not locking_parallel_barrier.task_exists():
                            child_ancestors_complete[ancestor] = False
                        elif locking_parallel_barrier.get_task_count() <= 0:
                            child_ancestors_complete[ancestor] = True
                    finally:
                        locking_parallel_barrier.release_lock()
                else:
                    base_scenario_ancestor_task_key = (
                        f"{base_case_execution_id}-{ancestor}"
                    )
                    locking_parallel_barrier = LockingParallelBarrier(
                        self.redis_url,
                        task_key=base_scenario_ancestor_task_key,
                        lock_key=f"{base_case_execution_id}-lock",
                    )
                    try:
                        locking_parallel_barrier.acquire_lock(
                            timeout=PARALLEL_PIPELINE_CALLBACK_BARRIER_TTL
                        )
                        if not locking_parallel_barrier.task_exists():
                            child_ancestors_complete[ancestor] = False
                        elif locking_parallel_barrier.get_task_count() <= 0:
                            child_ancestors_complete[ancestor] = True
                    finally:
                        locking_parallel_barrier.release_lock()

            if any(complete is False for complete in child_ancestors_complete.values()):
                logger.info(
                    "pypeline.child_ancestors_incomplete "
                    "pipeline_id=%s pipeline_name=%s actor_name=%s message_id=%s "
                    "execution_id=%s task_name=%s child=%s",
                    pipeline_id,
                    pipeline_name,
                    actor_name,
                    message_id,
                    execution_id,
                    task_name,
                    child,
                )
                continue

            if base_case_execution_id == execution_id:
                for scenario in message.options["scenarios"]:
                    child_ancestors = list(graph.predecessors(child))
                    child_has_other_ancestors_in_scenario = False

                    for ancestor in child_ancestors:
                        if ancestor in scenario["tasksToRunInScenario"]:
                            child_has_other_ancestors_in_scenario = True
                            break

                    if (
                        child in scenario["tasksToRunInScenario"]
                        and task_name in child_ancestors
                        and task_name not in scenario["tasksToRunInScenario"]
                        and not child_has_other_ancestors_in_scenario
                    ):
                        scenario_dispatch_key = self._child_dispatched_key(
                            base_case_execution_id, scenario["execution_id"], child
                        )
                        if not self.redis_conn.safe_set(
                            scenario_dispatch_key, "1", nx=True, ex=DEFAULT_RESULT_TTL
                        ):
                            logger.info(
                                "pypeline.scenario_child_already_dispatched "
                                "pipeline_id=%s pipeline_name=%s actor_name=%s message_id=%s "
                                "execution_id=%s task_name=%s child=%s scenario_execution_id=%s",
                                pipeline_id,
                                pipeline_name,
                                actor_name,
                                message_id,
                                execution_id,
                                task_name,
                                child,
                                scenario["execution_id"],
                            )
                            continue

                        task_key = f"{scenario['execution_id']}-{child}"
                        locking_parallel_barrier = LockingParallelBarrier(
                            self.redis_url,
                            task_key=task_key,
                            lock_key=f"{base_case_execution_id}-lock",
                        )
                        locking_parallel_barrier.set_task_count(1)
                        scenario_repl = scenario.get("taskReplacements") or {}
                        handler = task_definitions[child]["handlers"][
                            scenario_repl.get(child, 0)
                        ]
                        server_type = task_definitions[child].get("serverType", None)

                        lazy_actor = register_lazy_actor(
                            broker,
                            get_callable(handler),
                            pipeline_config["metadata"],
                            server_type,
                        )
                        scenario_message = lazy_actor.message()
                        scenario_message.options["pipeline"] = pipeline
                        if max_retries is not None:
                            scenario_message.options["max_retries"] = max_retries
                        scenario_message.options["task_replacements"] = scenario_repl
                        scenario_message.options["execution_id"] = scenario[
                            "execution_id"
                        ]

                        scenario_message.options["task_name"] = child
                        scenario_message.options["base_case_execution_id"] = (
                            message.options["base_case_execution_id"]
                        )
                        scenario_message.options["scenarios"] = message.options[
                            "scenarios"
                        ]
                        if "settings" in message.kwargs:
                            scenario_message.kwargs["settings"] = copy(
                                message.kwargs["settings"]
                            )
                            scenario_message.kwargs["settings"]["execution_id"] = (
                                scenario["execution_id"]
                            )

                        logger.info(
                            "pypeline.enqueue_scenario_child "
                            "pipeline_id=%s pipeline_name=%s actor_name=%s message_id=%s "
                            "execution_id=%s task_name=%s child=%s scenario_execution_id=%s",
                            pipeline_id,
                            pipeline_name,
                            actor_name,
                            message_id,
                            execution_id,
                            task_name,
                            child,
                            scenario["execution_id"],
                        )
                        messages.append(scenario_message)

            child_dispatch_key = self._child_dispatched_key(
                base_case_execution_id, execution_id, child
            )
            if not self.redis_conn.safe_set(
                child_dispatch_key, "1", nx=True, ex=DEFAULT_RESULT_TTL
            ):
                logger.info(
                    "pypeline.child_already_dispatched "
                    "pipeline_id=%s pipeline_name=%s actor_name=%s message_id=%s "
                    "execution_id=%s task_name=%s child=%s",
                    pipeline_id,
                    pipeline_name,
                    actor_name,
                    message_id,
                    execution_id,
                    task_name,
                    child,
                )
                continue

            task_key = f"{execution_id}-{child}"
            locking_parallel_barrier = LockingParallelBarrier(
                self.redis_url,
                task_key=task_key,
                lock_key=f"{base_case_execution_id}-lock",
            )
            locking_parallel_barrier.set_task_count(1)
            handler = task_definitions[child]["handlers"][
                task_replacements.get(child, 0)
            ]
            server_type = task_definitions[child].get("serverType", None)
            lazy_actor = register_lazy_actor(
                broker,
                get_callable(handler),
                pipeline_config["metadata"],
                server_type,
            )

            child_message = lazy_actor.message()
            child_message.options["pipeline"] = pipeline
            if max_retries is not None:
                child_message.options["max_retries"] = max_retries
            child_message.options["task_replacements"] = task_replacements
            child_message.options["execution_id"] = execution_id
            child_message.options["task_name"] = child
            child_message.options["base_case_execution_id"] = message.options[
                "base_case_execution_id"
            ]
            child_message.options["scenarios"] = message.options["scenarios"]
            if "settings" in message.kwargs:
                child_message.kwargs["settings"] = message.kwargs["settings"]
                child_message.kwargs["settings"]["execution_id"] = message.options[
                    "execution_id"
                ]

            logger.info(
                "pypeline.enqueue_child "
                "pipeline_id=%s pipeline_name=%s actor_name=%s message_id=%s "
                "execution_id=%s task_name=%s child=%s",
                pipeline_id,
                pipeline_name,
                actor_name,
                message_id,
                execution_id,
                task_name,
                child,
            )
            messages.append(child_message)

        for new_message in messages:
            broker.enqueue(new_message)
