import copy
import logging

from dramatiq.middleware 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

logger = logging.getLogger(__name__)


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

    def _get_execution_id(self, message):
        return message.kwargs.get("event", {}).get("execution_id")

    def _get_pipeline_id(self, message):
        return message.kwargs.get("event", {}).get("pipeline_id")

    def _get_pipeline_name(self, message):
        return message.kwargs.get("event", {}).get("pipeline_name")

    def _get_dispatching_key(self, execution_id: str, group_index: int) -> str:
        return f"pp:{execution_id}:group:{group_index}:dispatching"

    def _get_dispatched_key(self, execution_id: str, group_index: int) -> str:
        return f"pp:{execution_id}:group:{group_index}:dispatched"

    def after_process_message(self, broker, message, *, result=None, exception=None):
        from dramatiq.message import Message

        execution_id = self._get_execution_id(message)
        pipeline_id = self._get_pipeline_id(message)
        pipeline_name = self._get_pipeline_name(message)
        actor_name = getattr(message, "actor_name", None)
        message_id = getattr(message, "message_id", None)

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

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

            group_completion_uuid = message.options.get("group_completion_uuid")
            if not group_completion_uuid:
                logger.info(
                    "parallel_pipeline.after_process_message.skip_no_group_completion_uuid "
                    "pipeline_id=%s pipeline_name=%s execution_id=%s actor_name=%s message_id=%s",
                    pipeline_id,
                    pipeline_name,
                    execution_id,
                    actor_name,
                    message_id,
                )
                return

            execution_graph = message.options.get("execution_graph")
            if not execution_graph or not execution_id:
                logger.info(
                    "parallel_pipeline.after_process_message.skip_missing_execution_graph_or_execution_id "
                    "pipeline_id=%s pipeline_name=%s execution_id=%s actor_name=%s message_id=%s has_execution_graph=%s",
                    pipeline_id,
                    pipeline_name,
                    execution_id,
                    actor_name,
                    message_id,
                    bool(execution_graph),
                )
                return

            locking_parallel_barrier = LockingParallelBarrier(
                self.redis_url,
                task_key=group_completion_uuid,
                lock_key=f"{group_completion_uuid}-lock",
            )

            try:
                logger.info(
                    "parallel_pipeline.barrier.acquire "
                    "pipeline_id=%s pipeline_name=%s execution_id=%s actor_name=%s message_id=%s "
                    "group_completion_uuid=%s lock_key=%s timeout=%s",
                    pipeline_id,
                    pipeline_name,
                    execution_id,
                    actor_name,
                    message_id,
                    group_completion_uuid,
                    f"{group_completion_uuid}-lock",
                    PARALLEL_PIPELINE_CALLBACK_BARRIER_TTL,
                )
                locking_parallel_barrier.acquire_lock(
                    timeout=PARALLEL_PIPELINE_CALLBACK_BARRIER_TTL
                )
                remaining_tasks = locking_parallel_barrier.decrement_task_count()
                logger.info(
                    "parallel_pipeline.barrier.decremented "
                    "pipeline_id=%s pipeline_name=%s execution_id=%s actor_name=%s message_id=%s "
                    "group_completion_uuid=%s remaining_tasks=%s",
                    pipeline_id,
                    pipeline_name,
                    execution_id,
                    actor_name,
                    message_id,
                    group_completion_uuid,
                    remaining_tasks,
                )
            finally:
                released = locking_parallel_barrier.release_lock()
                logger.info(
                    "parallel_pipeline.barrier.released "
                    "pipeline_id=%s pipeline_name=%s execution_id=%s actor_name=%s message_id=%s "
                    "group_completion_uuid=%s released=%s",
                    pipeline_id,
                    pipeline_name,
                    execution_id,
                    actor_name,
                    message_id,
                    group_completion_uuid,
                    released,
                )

            if remaining_tasks > 0:
                logger.info(
                    "parallel_pipeline.after_process_message.waiting_on_group "
                    "pipeline_id=%s pipeline_name=%s execution_id=%s actor_name=%s message_id=%s "
                    "group_completion_uuid=%s remaining_tasks=%s",
                    pipeline_id,
                    pipeline_name,
                    execution_id,
                    actor_name,
                    message_id,
                    group_completion_uuid,
                    remaining_tasks,
                )
                return

            for i in range(len(execution_graph)):
                message_group = execution_graph[i]

                if (
                    message_group[0]["options"]["group_completion_uuid"]
                    != group_completion_uuid
                ):
                    continue

                logger.info(
                    "parallel_pipeline.group_completed "
                    "pipeline_id=%s pipeline_name=%s execution_id=%s actor_name=%s message_id=%s "
                    "completed_group_index=%s group_completion_uuid=%s",
                    pipeline_id,
                    pipeline_name,
                    execution_id,
                    actor_name,
                    message_id,
                    i,
                    group_completion_uuid,
                )

                if i + 1 >= len(execution_graph):
                    logger.info(
                        "parallel_pipeline.pipeline_complete "
                        "pipeline_id=%s pipeline_name=%s execution_id=%s actor_name=%s message_id=%s "
                        "completed_group_index=%s",
                        pipeline_id,
                        pipeline_name,
                        execution_id,
                        actor_name,
                        message_id,
                        i,
                    )
                    return

                next_group_index = i + 1
                next_group = execution_graph[next_group_index]

                dispatching_key = self._get_dispatching_key(
                    execution_id, next_group_index
                )
                dispatched_key = self._get_dispatched_key(
                    execution_id, next_group_index
                )

                if self.redis_conn.safe_exists(dispatched_key):
                    logger.info(
                        "parallel_pipeline.next_group_already_dispatched "
                        "pipeline_id=%s pipeline_name=%s execution_id=%s actor_name=%s message_id=%s "
                        "next_group_index=%s dispatched_key=%s",
                        pipeline_id,
                        pipeline_name,
                        execution_id,
                        actor_name,
                        message_id,
                        next_group_index,
                        dispatched_key,
                    )
                    return

                won_dispatch = self.redis_conn.safe_set(
                    dispatching_key,
                    "1",
                    nx=True,
                    ex=60,
                )
                if not won_dispatch:
                    logger.info(
                        "parallel_pipeline.next_group_dispatch_lock_not_acquired "
                        "pipeline_id=%s pipeline_name=%s execution_id=%s actor_name=%s message_id=%s "
                        "next_group_index=%s dispatching_key=%s",
                        pipeline_id,
                        pipeline_name,
                        execution_id,
                        actor_name,
                        message_id,
                        next_group_index,
                        dispatching_key,
                    )
                    return

                logger.info(
                    "parallel_pipeline.next_group_dispatch_lock_acquired "
                    "pipeline_id=%s pipeline_name=%s execution_id=%s actor_name=%s message_id=%s "
                    "next_group_index=%s dispatching_key=%s next_group_size=%s",
                    pipeline_id,
                    pipeline_name,
                    execution_id,
                    actor_name,
                    message_id,
                    next_group_index,
                    dispatching_key,
                    len(next_group),
                )

                try:
                    if self.redis_conn.safe_exists(dispatched_key):
                        logger.info(
                            "parallel_pipeline.next_group_already_dispatched_after_lock "
                            "pipeline_id=%s pipeline_name=%s execution_id=%s actor_name=%s message_id=%s "
                            "next_group_index=%s dispatched_key=%s",
                            pipeline_id,
                            pipeline_name,
                            execution_id,
                            actor_name,
                            message_id,
                            next_group_index,
                            dispatched_key,
                        )
                        return

                    completion_uuid = next_group[0]["options"]["group_completion_uuid"]
                    next_barrier = LockingParallelBarrier(
                        self.redis_url,
                        task_key=completion_uuid,
                        lock_key=f"{completion_uuid}-lock",
                    )
                    next_barrier.set_task_count(len(next_group))

                    logger.info(
                        "parallel_pipeline.next_group_barrier_initialized "
                        "pipeline_id=%s pipeline_name=%s execution_id=%s actor_name=%s message_id=%s "
                        "next_group_index=%s next_group_completion_uuid=%s next_group_size=%s",
                        pipeline_id,
                        pipeline_name,
                        execution_id,
                        actor_name,
                        message_id,
                        next_group_index,
                        completion_uuid,
                        len(next_group),
                    )

                    execution_graph_copy = copy.deepcopy(execution_graph)

                    for idx, next_message in enumerate(next_group):
                        next_message["options"][
                            "execution_graph"
                        ] = execution_graph_copy
                        logger.info(
                            "parallel_pipeline.enqueue_next_message "
                            "pipeline_id=%s pipeline_name=%s execution_id=%s actor_name=%s message_id=%s "
                            "next_group_index=%s next_message_index=%s next_actor_name=%s next_group_completion_uuid=%s",
                            pipeline_id,
                            pipeline_name,
                            execution_id,
                            actor_name,
                            message_id,
                            next_group_index,
                            idx,
                            next_message.get("actor_name"),
                            completion_uuid,
                        )
                        broker.enqueue(Message(**next_message))

                    self.redis_conn.safe_set(
                        dispatched_key,
                        "1",
                        ex=DEFAULT_RESULT_TTL,
                    )

                    logger.info(
                        "parallel_pipeline.next_group_dispatched "
                        "pipeline_id=%s pipeline_name=%s execution_id=%s actor_name=%s message_id=%s "
                        "next_group_index=%s dispatched_key=%s",
                        pipeline_id,
                        pipeline_name,
                        execution_id,
                        actor_name,
                        message_id,
                        next_group_index,
                        dispatched_key,
                    )
                finally:
                    self.redis_conn.safe_delete(dispatching_key)
                    logger.info(
                        "parallel_pipeline.next_group_dispatch_lock_released "
                        "pipeline_id=%s pipeline_name=%s execution_id=%s actor_name=%s message_id=%s "
                        "next_group_index=%s dispatching_key=%s",
                        pipeline_id,
                        pipeline_name,
                        execution_id,
                        actor_name,
                        message_id,
                        next_group_index,
                        dispatching_key,
                    )

                return

            logger.warning(
                "parallel_pipeline.completed_group_not_found_in_execution_graph "
                "pipeline_id=%s pipeline_name=%s execution_id=%s actor_name=%s message_id=%s "
                "group_completion_uuid=%s",
                pipeline_id,
                pipeline_name,
                execution_id,
                actor_name,
                message_id,
                group_completion_uuid,
            )

        except Exception:
            logger.exception(
                "parallel_pipeline.after_process_message.failed "
                "pipeline_id=%s pipeline_name=%s execution_id=%s actor_name=%s message_id=%s "
                "group_completion_uuid=%s",
                pipeline_id,
                pipeline_name,
                execution_id,
                actor_name,
                message_id,
                message.options.get("group_completion_uuid"),
            )
            raise
