from __future__ import annotations

import copy
import json
import time
import typing
from uuid import uuid4

from dramatiq.broker import get_broker
from dramatiq.results import ResultMissing
from redis.exceptions import RedisError

from pypeline.barrier import LockingParallelBarrier
from pypeline.constants import (
    REDIS_URL,
    DEFAULT_RESULT_TTL,
)
from pypeline.utils.redis_client import ResilientRedisClient

from dramatiq.message import Message


class parallel_pipeline:
    """Chain actors together, passing the result of one actor to the
    next one in line.

    Parameters:
      children(typing.List[typing.List[Message]]): A sequence of messages or
        pipelines. Child pipelines are flattened into the resulting
        pipeline.
      broker(Broker): The broker to run the pipeline on. Defaults to
        the current global broker.
    """

    messages: list[Message]

    def __init__(self, messages: typing.List[typing.List[Message]], broker=None):
        self.broker = broker or get_broker()
        self.messages = messages
        self.execution_id = str(uuid4())
        execution_graph = []

        for group_index, message_group in enumerate(self.messages):
            sub_execution_group = []
            group_completion_uuid = str(uuid4())
            for task_index, m in enumerate(message_group):
                m.kwargs["event"]["execution_id"] = self.execution_id
                m.options["group_completion_uuid"] = group_completion_uuid
                m.options["group_index"] = group_index
                m.options["task_index"] = task_index
                message_dict = copy.deepcopy(m.asdict())
                sub_execution_group.append(message_dict)
            execution_graph.append(sub_execution_group)

        self.execution_graph = execution_graph

        for m in self.messages[0]:
            m.options["execution_graph"] = execution_graph

    def __len__(self):
        """Returns the length of the parallel_pipeline."""
        count = 0
        for message_group in self.messages:
            count += len(message_group)
        return count

    def __str__(self):  # pragma: no cover
        """Return a string representation of the parallel_pipeline."""
        result = []

        for i, message_group in enumerate(self.messages):
            group_str = f"Group {i + 1}: [\n"
            for j, message in enumerate(message_group):
                group_str += f"  Message {j + 1}: {message.actor_name}\n"
            group_str += "]\n"
            result.append(group_str)

        return "".join(result)

    @property
    def completed(self):
        return self.completed_count == len(self)

    @property
    def completed_count(self):
        count = 0

        for message_group in self.messages:
            for message in message_group:
                try:
                    message.get_result()
                    count += 1
                except ResultMissing:
                    pass
        return count

    def run(self, *, delay=None):
        """Run this parallel_pipeline."""
        starting_group = self.messages[0]

        completion_uuid = starting_group[0].options["group_completion_uuid"]
        locking_parallel_barrier = LockingParallelBarrier(
            REDIS_URL, task_key=completion_uuid, lock_key=f"{completion_uuid}-lock"
        )
        locking_parallel_barrier.set_task_count(len(starting_group))

        redis_conn = ResilientRedisClient(REDIS_URL)
        redis_conn.safe_set(
            self._group_dispatched_key(self.execution_id, 0),
            "1",
            ex=DEFAULT_RESULT_TTL,
        )

        for m in starting_group:
            self.broker.enqueue(m, delay=delay)

        return self

    def get_result(self, *, block=False, timeout=None):
        """Get the result of this pipeline."""
        last_message = self.messages[-1][-1]
        backend = self.broker.get_results_backend()
        return last_message.get_result(backend=backend, block=block, timeout=timeout)

    def get_results(self, *, block=False, timeout=None):
        """Get the results of each job in the pipeline."""
        deadline = None
        if timeout:
            deadline = time.monotonic() + timeout / 1000

        for message_group in self.messages:
            for message in message_group:
                if deadline:
                    timeout = max(0, int((deadline - time.monotonic()) * 1000))

                backend = self.broker.get_results_backend()
                yield {
                    message.actor_name: message.get_result(
                        backend=backend, block=block, timeout=timeout
                    )
                }

    def to_json(self) -> str:
        """Convert the execution graph to a JSON string representation."""
        return json.dumps(self.execution_graph)

    @classmethod
    def from_json(cls, json_data: str) -> parallel_pipeline:
        """Reconstruct a parallel_pipeline from serialized execution graph."""
        execution_graph = json.loads(json_data)

        obj = cls.__new__(cls)
        obj.broker = get_broker()
        obj.execution_graph = execution_graph
        obj.messages = [
            [Message(**message_dict) for message_dict in message_group]
            for message_group in execution_graph
        ]

        try:
            obj.execution_id = obj.messages[0][0].kwargs["event"]["execution_id"]
        except (IndexError, KeyError):
            obj.execution_id = None

        return obj

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


class PipelineResult:
    """
    A class to manage and retrieve the results of a parallel pipeline execution.
    """

    def __init__(self, execution_id: str, result_ttl: int = DEFAULT_RESULT_TTL):
        self.pipeline: parallel_pipeline = None
        self.execution_id = execution_id
        self.redis_key = f"{execution_id}-results-key"
        self.result_ttl = result_ttl
        self.redis_conn = ResilientRedisClient(REDIS_URL)

    def create_result_entry(self, pipeline_json_str: str):
        if not pipeline_json_str:
            raise ValueError("No pipeline data passed to create result store")

        try:
            self.redis_conn.safe_setex(
                self.redis_key, self.result_ttl, pipeline_json_str
            )
        except RedisError as e:
            raise RuntimeError(f"Failed to store pipeline data in Redis: {e}")

    def load(self):
        try:
            pipeline_data = self.redis_conn.safe_get(self.redis_key)
            if pipeline_data:
                if isinstance(pipeline_data, bytes):
                    pipeline_data = pipeline_data.decode("utf-8")
                self.pipeline = parallel_pipeline.from_json(pipeline_data)
            else:
                self.pipeline = None
        except RedisError as e:
            raise RuntimeError(f"Failed to load pipeline data from Redis: {e}")

    @property
    def status(self) -> str:
        if not self.pipeline:
            return "unavailable"
        return "complete" if self.pipeline.completed else "pending"

    def get_results(self) -> dict:
        if not self.pipeline:
            return {}

        results = {}
        for result in self.pipeline.get_results():
            for actor, res in result.items():
                unique_actor = self._get_unique_actor_name(actor, results)
                results[unique_actor] = res
        return results

    def get_result(self):
        if self.pipeline:
            return self.pipeline.get_result()

    def _get_unique_actor_name(self, actor: str, results: dict) -> str:
        if actor not in results:
            return actor

        suffix = 0
        new_actor = f"{actor}-{suffix}"
        while new_actor in results:
            suffix += 1
            new_actor = f"{actor}-{suffix}"
        return new_actor
