from __future__ import annotations

import copy
import json
import time
import typing
from uuid import uuid4
from urllib.parse import urlparse

from dramatiq.broker import get_broker
from dramatiq.results import ResultMissing
import redis
from redis.backoff import ExponentialBackoff
from redis.exceptions import RedisError, ConnectionError, TimeoutError, BusyLoadingError, TryAgainError
from redis.retry import Retry
from redis.sentinel import Sentinel, MasterNotFoundError

from pypeline.barrier import LockingParallelBarrier
from pypeline.constants import (
    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 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))

        # Mark group 0 as already dispatched so later logic is consistent.
        redis_conn = self._get_redis_connection()
        redis_conn.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 _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 _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

        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,
            )
            self.redis_conn = 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,
            )
        else:
            self.redis_conn = redis.StrictRedis.from_url(REDIS_URL, **common_kwargs)

    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.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.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
