import logging
from copy import copy
from urllib.parse import urlparse

import networkx as nx
import redis
from dramatiq import Middleware
from redis.backoff import ExponentialBackoff
from redis.retry import Retry
from redis.sentinel import Sentinel, MasterNotFoundError
from redis.exceptions import (
    ConnectionError,
    TimeoutError,
    BusyLoadingError,
    TryAgainError,
)

from pypeline.barrier import LockingParallelBarrier
from pypeline.constants import (
    PARALLEL_PIPELINE_CALLBACK_BARRIER_TTL,
    REDIS_URL,
    REDIS_SENTINEL_MASTER_NAME,
    DEFAULT_REDIS_SOCKET_CONNECT_TIMEOUT,
    DEFAULT_REDIS_SOCKET_TIMEOUT,
    DEFAULT_REDIS_RETRY_ON_TIMEOUT,
    DEFAULT_REDIS_SOCKET_KEEPALIVE,
    DEFAULT_REDIS_HEALTH_CHECK_INTERVAL,
    DEFAULT_RESULT_TTL,
)
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 = self._get_redis_connection()

    def _get_redis_connection(self):
        retry = Retry(ExponentialBackoff(), 5)

        common_kwargs = {
            "socket_connect_timeout": DEFAULT_REDIS_SOCKET_CONNECT_TIMEOUT,
            "socket_timeout": DEFAULT_REDIS_SOCKET_TIMEOUT,
            "retry_on_timeout": DEFAULT_REDIS_RETRY_ON_TIMEOUT,
            "socket_keepalive": DEFAULT_REDIS_SOCKET_KEEPALIVE,
            "health_check_interval": DEFAULT_REDIS_HEALTH_CHECK_INTERVAL,
            "retry": retry,
            "retry_on_error": [
                ConnectionError,
                TimeoutError,
                MasterNotFoundError,
                BusyLoadingError,
                TryAgainError,
            ],
        }

        if REDIS_SENTINEL_MASTER_NAME is not None:
            parsed_redis_url = urlparse(REDIS_URL)
            redis_sentinel = Sentinel(
                sentinels=[(parsed_redis_url.hostname, parsed_redis_url.port)],
                sentinel_kwargs=common_kwargs,
            )
            return redis_sentinel.master_for(
                REDIS_SENTINEL_MASTER_NAME,
                db=int(parsed_redis_url.path[1]) if parsed_redis_url.path else 0,
                password=parsed_redis_url.password,
                **common_kwargs,
            )

        return redis.StrictRedis.from_url(
            REDIS_URL,
            **common_kwargs,
        )

    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,
        )

        # Signal to other jobs that current task is finished
        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=10)
            _ = 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.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)

            # If we've made it here all ancestors of this child are complete, and it's time to run.
            # Use a dispatch gate to prevent double-dispatch when multiple ancestors complete simultaneously.
            child_dispatch_key = self._child_dispatched_key(
                base_case_execution_id, execution_id, child
            )
            if not self.redis_conn.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)
