import asyncio
from contextlib import asynccontextmanager
import contextlib
from functools import cached_property, lru_cache
import inspect
import logging
import tenacity
import pydantic
from typing import (
    TYPE_CHECKING,
    Any,
    Final,
    Generic,
    Literal,
    TypeVar,
    TypedDict,
    cast,
)
from collections.abc import AsyncIterator, Awaitable, Callable, Iterator
from celery.exceptions import TaskRevokedError
from celery.result import AsyncResult
from celery import Signature, states
from acelery.errors import MaxPendingInterruptsError

from redis.asyncio.client import PubSub
from celery.backends.base import Backend
from celery.backends.redis import RedisBackend
from redis import asyncio as aioredis
from redis.asyncio.connection import ConnectionPool
from typing_extensions import NotRequired

from acelery.executable import AIOExecutable


if (major := pydantic.__version__.split(".", 1)[0]) == "1":
    from pydantic import parse_obj_as
elif major == "2":
    from pydantic import TypeAdapter  # type: ignore[attr-defined]

    _T = TypeVar("_T")
    NameFactory = str | Callable[[type[Any]], str]

    def parse_obj_as(
        type_: type[_T], obj: Any, *, type_name: NameFactory | None = None
    ) -> _T:
        return TypeAdapter(type_).validate_python(obj)
else:
    raise ValueError(f"Unsupported pydantic version: {pydantic.__version__}")


logger = logging.getLogger(__name__)


if TYPE_CHECKING:
    # for some reason the IDE & mypy can't recognize the TypeForm type although it does exist in typing_extensions
    # TypeForm is the correct way to type hint a type of class or union type and is fully available in python >3.14
    # the type does exist in the typing_extensions module, but the IDE & mypy doesn't recognize it
    # see https://github.com/python/typing_extensions/blob/main/src/typing_extensions.py#L2091
    from typing_extensions import TypeForm  # type: ignore[attr-defined]

T = TypeVar("T")

LUA_DEL_IF_OWNED = """
if redis.call('GET', KEYS[1]) == ARGV[1] then
    return redis.call('DEL', KEYS[1])
end
return 0
"""


class CeleryResult(TypedDict):
    status: str
    result: Any | None
    traceback: NotRequired[str | None]
    queue: NotRequired[str]


class AIOResult(Generic[T]):
    """
    A wrapper around the celery AsyncResult that provides an asyncio safe way to
    wait for the result of the task
    and to trigger or join a tasks uniquely based on their workflow, args and kwargs.
    """

    DEFAULT_FORGET_COOLDOWN: Final[int] = 90
    DEFAULT_INTERRUPT_EVERY: Final[float] = 5.0
    DEFAULT_MAX_PENDING_INTERRUPTS: Final[int] = 3

    def __init__(
        self, celery_result: AsyncResult, return_type: "type[T] | TypeForm[T]"
    ):
        self.celery_result = celery_result
        self._backend: RedisBackend = celery_result.app.backend
        self.redis = self.extract_redis_from_backend(self._backend)
        self.return_type = return_type
        self._task_meta: CeleryResult = {"status": states.PENDING, "result": None}

    def with_type(self, return_type: "type[T] | TypeForm[T]") -> "AIOResult[T]":
        """
        Returns a new AsyncioResult with the given return type

        Args:
            return_type: The return type of the task

        Returns:
            A new AsyncioResult with the given return type
        """
        asyncio_result: AIOResult[T] = AIOResult(self.celery_result, return_type)
        asyncio_result._backend = self._backend
        asyncio_result._task_meta = self._task_meta
        asyncio_result.redis = self.redis
        return asyncio_result

    @staticmethod
    @lru_cache(maxsize=10)
    def _extract_redis_from_backend(
        backend: Backend, _: asyncio.AbstractEventLoop
    ) -> aioredis.Redis:
        """
        Extracts the redis client from the celery result backend, this is cached
        by the backend and the event loop objects to avoid creating new redis clients
        for every call


        Args:
            backend: The celery result backend
            _: The event loop object, this is required by the lru_cache decorator

        Returns:
            The redis client
        """
        if backend is None or not isinstance(backend, RedisBackend):
            raise ValueError("Celery result backend is not a Redis backend")
        conn_pool = ConnectionPool(**backend.connparams)
        redis = aioredis.Redis(connection_pool=conn_pool)
        redis.auto_close_connection_pool = True
        return redis

    @staticmethod
    def extract_redis_from_backend(backend: Backend) -> aioredis.Redis:
        """
        Extracts the redis client from the celery result backend

        Args:
            backend: The celery result backend

        Returns:
            The redis client
        """
        return AIOResult._extract_redis_from_backend(
            backend, asyncio.get_running_loop()
        )

    @cached_property
    def task_id(self) -> str:
        """
        Gets the task id from the celery result

        Returns:
            The task id
        """
        return self.celery_result.id

    @cached_property
    def task_key(self) -> str:
        """
        Gets the task key from the celery result, this is the key in the redis backend

        Returns:
            The task key
        """
        return f"{str(self._backend.task_keyprefix, encoding='utf-8')}{self.celery_result.id}"

    @property
    def task_status(self) -> str:
        """
        The latest known status of the task

        Returns:
            The latest known status of the task
        """
        return self._task_meta["status"]

    @asynccontextmanager
    async def _pubsub_ctx(self, channel: str) -> AsyncIterator[PubSub]:
        """
        Creates a context manager for the redis pubsub that subscribes on enter and unsubscribes on exit

        Args:
            channel: The channel to subscribe to

        Yields:
            The redis pubsub object
        """
        async with self.redis.pubsub() as pubsub:
            await pubsub.subscribe(channel)
            try:
                yield pubsub
            finally:
                await pubsub.unsubscribe(channel)

    async def _get_task_meta(
        self,
        *,
        missing_level: Literal["DEBUG", "INFO", "WARNING", "ERROR"] = "WARNING",
        new_expiry: int | None = None,
    ) -> CeleryResult:
        """
        Gets the task meta from the redis backend

        Args:
            missing_level: The level to log if the task meta is not found in the redis backend
            new_expiry: The new expiry to set for the task meta in the redis backend, None means no modification

        Returns:
            The task meta
        """
        if self.celery_result._cache is not None:  # type: ignore [attr-defined]
            return self.celery_result._cache  # type: ignore [attr-defined]
        if new_expiry is not None:
            task_meta = await self.redis.getex(self.task_key, ex=new_expiry)
        else:
            task_meta = await self.redis.get(self.task_key)
        if task_meta is None:
            logger.log(
                logging._nameToLevel[missing_level],
                "Task %s not found in redis",
                self.task_id,
            )
            return self._task_meta
        self._task_meta = self._backend.meta_from_decoded(
            self._backend.decode_result(task_meta)
        )
        return self._task_meta

    def _extract_result_from_task_meta(
        self, task_meta: CeleryResult, *, propagate: bool
    ) -> tuple[Any | None, bool, bool]:
        """
        Extracts the result from the decoded tasks meta

        Args:
            task_meta: The decoded tasks meta
            propagate: Whether to propagate the exception

        Returns:
            A tuple containing the result, whether the task is ready, and whether the task is pending
        """
        _status, _result = task_meta["status"], task_meta["result"]
        logger.debug(
            "Extracting result from task meta of %s with status %s",
            self.task_id,
            _status,
        )
        if _status == states.PENDING:
            return None, False, True
        elif _status in states.READY_STATES:
            logger.debug(
                "Setting task meta cache in celery result for ready task %s",
                self.task_id,
            )
            self.celery_result._set_cache(task_meta)  # type: ignore [attr-defined]
            # if the task is ready, we either return the result or throw the exception
            if _status == states.SUCCESS:
                logger.debug(
                    "Task %s is ready, returning result",
                    self.task_id,
                )
                return _result, True, False
            self.celery_result.maybe_throw(propagate=propagate)
            logger.debug(
                "Task %s is ready, returning result",
                self.task_id,
            )
            return _result, True, False
        else:
            logger.debug(
                "Task %s is not ready yet",
                self.task_id,
            )
            return None, False, False

    async def _handle_pubsub_message(
        self,
        message: dict,
        propagate: bool,
        channels: list[bytes],
    ) -> tuple[dict | None, bool, bool]:
        """
        Handles the pubsub message from the redis pubsub

        Args:
            message: The raw message from the redis pubsub
            propagate: Whether to propagate the exception
            channels: The channels from the redis pubsub

        Returns:
            A tuple containing the result, whether the task is ready, and whether the task is pending
        """
        # an inline fnction so we can use wait_for for timeout
        if message.get("type") == "subscribe":
            logger.debug(
                "Received subscribe message from redis pubsub for %s",
                self.task_id,
            )
            task_meta = await self._get_task_meta(missing_level="DEBUG")
            return self._extract_result_from_task_meta(task_meta, propagate=propagate)

        logger.debug(
            "Received data from pubsub channels %s",
            ", ".join(map(lambda c: c.decode("utf-8"), channels)),
        )
        raw_data = message.get("data")
        if raw_data is None:
            logger.warning(
                "Message data is None, "
                "falling back to extraction result from task meta stored in redis"
            )
            task_meta = await self._get_task_meta()
            return self._extract_result_from_task_meta(task_meta, propagate=propagate)
        self._task_meta = self._backend.meta_from_decoded(
            self._backend.decode_result(raw_data)
        )
        return self._extract_result_from_task_meta(self._task_meta, propagate=propagate)

    async def _handle_tick(
        self,
        propagate: bool,
    ) -> tuple[dict | None, bool, bool]:
        """
        Handles the tick message from the redis pubsub
        This is used to check the task meta in case we haven't gotten any message from
        the pubsub listener for a while

        Args:
            propagate: Whether to propagate the exception

        Returns:
            A tuple containing the result, whether the task is ready, and whether the task is pending
        """
        logger.debug(
            "Received tick for task %s",
            self.task_id,
        )
        task_meta = await self._get_task_meta()
        return self._extract_result_from_task_meta(task_meta, propagate=propagate)

    async def _await_result(
        self,
        pubsub: PubSub,
        *,
        propagate: bool,
        interrupt_every: float,
        max_pending_interrupts: int,
    ) -> Any:
        """
        Awaits the result of the task

        Args:
            pubsub: The pubsub object
            propagate: Whether to propagate the exception
            interrupt_every: The interval to check the task meta in case we haven't gotten any message from
            the pubsub listener for a while
            max_pending_interrupts: The maximum number of pending interrupts, practically saying a timeout of
             max_pending_interrupts * interrupt_every for the task to be received in the worker

        Returns:
            The raw result of the task
        """
        pending_interrupts = 0
        channels = list(pubsub.channels.keys())
        pubsub_listener: AsyncIterator[dict] = pubsub.listen()
        listener_task: asyncio.Task[dict] = asyncio.create_task(anext(pubsub_listener))  # type: ignore[arg-type]
        tick_task = asyncio.create_task(asyncio.sleep(interrupt_every))
        try:
            while True:
                # to prevent any race conditions of never getting the subscribe message
                # of the task finishing, we wait for the first task to complete
                # between
                # - getting the next subscribe message
                # - sleeping for the interrupt_every interval
                # if the listener task completes, we handle the message
                # if the tick task completes, we handle the tick and get the task from the redis
                done, _ = await asyncio.wait(
                    [listener_task, tick_task],
                    return_when=asyncio.FIRST_COMPLETED,
                )
                if listener_task in done:
                    message = listener_task.result()
                    result, is_ready, is_pending = await self._handle_pubsub_message(
                        message, propagate, channels
                    )
                    if is_ready:
                        # the finally block will cancel the listener task
                        return result
                    else:
                        if is_pending:
                            pending_interrupts += 1
                            if pending_interrupts > max_pending_interrupts:
                                logger.warning(
                                    "Max pending interrupts of %s exceeded for task %s",
                                    max_pending_interrupts,
                                    self.task_id,
                                )
                                raise MaxPendingInterruptsError(
                                    self.task_id, max_pending_interrupts
                                )
                        tick_task.cancel()
                        listener_task = asyncio.create_task(anext(pubsub_listener))  # type: ignore[arg-type]
                        # we create a new tick task to avoid race conditions
                        # to wait for the next tick message starting from now
                        tick_task = asyncio.create_task(asyncio.sleep(interrupt_every))
                        continue
                else:
                    result, is_ready, is_pending = await self._handle_tick(propagate)
                    if is_ready:
                        # the finally block will cancel the listener task
                        return result
                    else:
                        if is_pending:
                            pending_interrupts += 1
                            if pending_interrupts > max_pending_interrupts:
                                logger.warning(
                                    "Max pending interrupts of %s reached for task %s",
                                    max_pending_interrupts,
                                    self.task_id,
                                )
                                raise MaxPendingInterruptsError(
                                    self.task_id, max_pending_interrupts
                                )
                        # we create a new tick task to avoid race conditions
                        # to wait for the next tick message starting from now
                        tick_task = asyncio.create_task(asyncio.sleep(interrupt_every))
                        continue
        except asyncio.CancelledError:
            raise
        except Exception:
            logger.exception(
                "Error awaiting result for task %s", self.task_id, exc_info=True
            )
            raise
        finally:
            listener_task.cancel()
            tick_task.cancel()
            try:
                if inspect.isasyncgen(pubsub_listener):
                    with contextlib.suppress(
                        asyncio.CancelledError, asyncio.TimeoutError
                    ):
                        await asyncio.wait_for(listener_task, timeout=5)
                    with contextlib.suppress(GeneratorExit, asyncio.TimeoutError):
                        await asyncio.wait_for(pubsub_listener.aclose(), timeout=5)
                else:
                    logger.warning(
                        "Pubsub listener is not a generator for task %s, this should never happen",
                        self.task_id,
                    )
            except Exception:
                logger.exception(
                    "Error closing pubsub listener for task %s",
                    self.task_id,
                )
                raise
            else:
                logger.debug(
                    "Pubsub listener generator closed for task %s",
                    self.task_id,
                )

    async def get(
        self,
        timeout: float = 60,
        *,
        propagate: bool = True,
        forget_cooldown: int = DEFAULT_FORGET_COOLDOWN,
        forget: bool = True,
        interrupt_every: float = DEFAULT_INTERRUPT_EVERY,
        max_pending_interrupts: int = DEFAULT_MAX_PENDING_INTERRUPTS,
    ) -> T:
        """
        Gets the result of the task

        Args:
            timeout: The timeout to wait for the task to finish
            propagate: Whether to propagate the exception
            forget_cooldown: The cooldown to forget the task result
            forget: Whether to forget the task result
            interrupt_every: The interval to check the task meta in case we haven't gotten any message from
            the pubsub listener for a while
            max_pending_interrupts: The maximum number of pending interrupts, practically saying a timeout of
            max_pending_interrupts * interrupt_every for the task to be received in the worker

        Returns:
            The parsed result of the task
        """
        if interrupt_every > timeout:
            previous_interrupt_every = interrupt_every
            # the race here is that the task was finished before we started to subscribe
            # causing us to never get the subscribe message and timeout
            interrupt_every = timeout / max_pending_interrupts
            logger.warning(
                "Attempted to set interrupt every to %s seconds, "
                "which is greater than the timeout of %s seconds. "
                "Reduced the interrupt every to %s(%s/%s == %s/%s) seconds to prevent "
                "race conditions",
                previous_interrupt_every,
                interrupt_every,
                max_pending_interrupts,
                timeout,
                max_pending_interrupts,
            )

        logger.info(
            "Waiting for result of task %s with timeout %s",
            self.task_id,
            timeout,
        )
        async with self._pubsub_ctx(self.task_key) as pubsub:
            logger.debug("Subscribed to pubsub channel %s", self.task_key)
            try:
                exc: Exception | None = None
                result = await asyncio.wait_for(
                    self._await_result(
                        pubsub,
                        propagate=propagate,
                        interrupt_every=interrupt_every,
                        max_pending_interrupts=max_pending_interrupts,
                    ),
                    timeout=timeout,
                )
            except asyncio.TimeoutError as e:
                exc = e
                logger.exception(
                    "Timeout of %s seconds waiting for workflow %s to finish, last known status is %s",
                    timeout,
                    self.task_id,
                    self.task_status,
                )
                raise TimeoutError(
                    f"Timeout of {timeout} seconds waiting for workflow {self.task_id} "
                    f"to finish, last known status is {self.task_status}"
                ) from e
            except Exception as e:
                exc = e
                raise
            else:
                logger.debug(
                    "Got result for task %s, parsing result",
                    self.task_id,
                )
                return parse_obj_as(self.return_type, result)
            finally:
                if forget and exc is None:
                    logger.debug(
                        "forgetting task %s result with cooldown of %s",
                        self.task_id,
                        forget_cooldown,
                    )
                    # we set the expiry of the task meta key to a configurable seconds from now to allow other possible clients
                    # to get the result.
                    # This is similar to calling celery_result.forget with the same cooldown
                    # we only extend the expiry if the new expiry is greater than the current expiry
                    await self.redis.expire(self.task_key, forget_cooldown)

    @staticmethod
    def _convert_workflow_to_unique_running_task_key(workflow: Signature) -> str:
        """
        Converts the workflow to a unique running task key based on the workflow name, args and kwargs.
        Two tasks with the same name, args and kwargs will have the same unique running task key.

        Args:
            workflow: The workflow to convert

        Returns:
            The unique running task key
        """

        def _repr_iter() -> Iterator[str]:
            if not workflow.args and not workflow.kwargs:
                yield ""
                return
            if workflow.args:
                yield from map(lambda arg: f"{repr(arg)}", workflow.args)
            if workflow.kwargs:
                for k, v in workflow.kwargs.items():
                    yield f"{repr(k)}={repr(v)}"

        workflow_args = (
            f"{','.join(_repr_iter())}" if workflow.args or workflow.kwargs else ""
        )

        return f"{workflow.name}({workflow_args})"

    @classmethod
    async def _publish_new_task(
        cls,
        workflow: Signature,
        *,
        timeout: float,
        return_type: "type[T] | TypeForm[T]",
        unique_running_task_key: str,
        forget_cooldown: int,
    ) -> "AIOResult[T]":
        """
        Publishes a new task with the given workflow, timeout, return type and
        sets the newly published task id in the redis backend with the given
        unique running task key and forget cooldown.

        Args:
            workflow: The workflow to publish
            timeout: The timeout to wait for the task to finish
            return_type: The return type of the task
            unique_running_task_key: The unique running task key
            forget_cooldown: The cooldown to forget the task result

        Returns:
            The AsyncioResult of the newly published task
        """
        async_result = await AIOExecutable(
            workflow, return_type=return_type
        ).apply_async(expires=timeout)
        await async_result.redis.set(
            unique_running_task_key,
            async_result.task_id,
            ex=int(timeout + forget_cooldown),
        )
        return async_result

    @classmethod
    async def _trigger_or_join_task(
        cls,
        workflow: Signature,
        *,
        timeout: float,
        return_type: "type[T] | TypeForm[T]",
        propagate: bool,
        forget_cooldown: int,
        forget: bool,
        interrupt_every: float,
        max_pending_interrupts: int,
        track_started: bool,
    ) -> tuple["AIOResult[T]", T]:
        """
        Triggers or joins a task based on the unique running task key.

        Args:
            workflow: The workflow to trigger or join
            timeout: The timeout to wait for the task to finish
            return_type: The return type of the task
            propagate: Whether to propagate the exception
            forget_cooldown: The cooldown to forget the task result
            forget: Whether to forget the task result
            interrupt_every: The interval to check the task meta in case we haven't gotten any message from
            the pubsub listener for a while
            max_pending_interrupts: The maximum number of pending interrupts, practically saying a timeout of
            max_pending_interrupts * interrupt_every for the task to be received in the worker

        Returns:
            A tuple containing the AsyncioResult and the result of the task
        """
        unique_running_task_key = cls._convert_workflow_to_unique_running_task_key(
            workflow
        )
        redis = cls.extract_redis_from_backend(workflow.app.backend)
        task_id_bytes = cast(bytes | None, await redis.get(unique_running_task_key))
        task_id = task_id_bytes.decode("utf-8") if task_id_bytes is not None else None
        if task_id is None:
            async_result = await cls._publish_new_task(
                workflow,
                timeout=timeout,
                return_type=return_type,
                unique_running_task_key=unique_running_task_key,
                forget_cooldown=forget_cooldown,
            )
            logger.debug(
                "Could not find running task with key %s, "
                "published a new one with id %s...",
                unique_running_task_key,
                async_result.task_id,
            )
        else:
            async_result = cls(AsyncResult(task_id, app=workflow.app), return_type)
            task_meta = await async_result._get_task_meta(
                new_expiry=int(timeout + forget_cooldown)
            )
            if (state := task_meta.get("status")) in states.READY_STATES or (
                track_started and state == states.PENDING
            ):
                async_result = await cls._publish_new_task(
                    workflow,
                    timeout=timeout,
                    return_type=return_type,
                    unique_running_task_key=unique_running_task_key,
                    forget_cooldown=forget_cooldown,
                )
                logger.debug(
                    "Found existing task for key %s with id %s, "
                    "but it's state is %s. "
                    "Published a new one with id %s",
                    unique_running_task_key,
                    task_id,
                    state,
                    async_result.task_id,
                )
            else:
                logger.info(
                    "Found existing running task for key %s with id %s, "
                    "joining on the result",
                    unique_running_task_key,
                    async_result.task_id,
                )
        try:
            result = await async_result.get(
                timeout=timeout,
                propagate=propagate,
                forget_cooldown=forget_cooldown,
                forget=forget,
                interrupt_every=interrupt_every,
                max_pending_interrupts=max_pending_interrupts,
            )
        finally:
            is_deleted = int(
                await cast(
                    Awaitable[str],
                    redis.eval(
                        LUA_DEL_IF_OWNED,
                        1,
                        unique_running_task_key,
                        async_result.task_id,
                    ),
                )
            )
            if is_deleted == 1:
                logger.debug(
                    "Invalidated unique running task key %s with task id %s",
                    unique_running_task_key,
                    async_result.task_id,
                )
            else:
                logger.debug(
                    "Cached unique running task in key %s is not owned by task %s, skipped invalidation",
                    unique_running_task_key,
                    async_result.task_id,
                )

        return async_result, result

    @classmethod
    async def trigger_or_join_task(
        cls,
        workflow: Signature,
        *,
        timeout: float,
        return_type: "type[T] | TypeForm[T]",
        propagate: bool = True,
        forget_cooldown: int = DEFAULT_FORGET_COOLDOWN,
        forget: bool = True,
        interrupt_every: float = DEFAULT_INTERRUPT_EVERY,
        retry_on_timeout: bool = True,
        max_retries_on_timeout: int = 3,
        max_pending_interrupts: int = DEFAULT_MAX_PENDING_INTERRUPTS,
        track_started: bool = False,
    ) -> tuple["AIOResult[T]", T]:
        """
        Triggers or joins a task based on the unique running task key with configurable
        retries on timeouts

        Args:
            workflow: The workflow to trigger or join
            timeout: The timeout to wait for the task to finish
            return_type: The return type of the task
            propagate: Whether to propagate the exception
            forget_cooldown: The cooldown to forget the task result
            forget: Whether to forget the task result
            interrupt_every: The interval to check the task meta in case we haven't gotten any message from
            the pubsub listener for a while
            retry_on_timeout: Whether to retry on timeout
            max_retries_on_timeout: The maximum number of retries on timeout
            max_pending_interrupts: The maximum number of pending interrupts, practically saying a timeout of
            max_pending_interrupts * interrupt_every for the task to be received in the worker
            track_started: Whether the task is configured to track started, if True, then
             the trigger will consider PENDING state as not started instead of potentially started task

        Returns:
            A tuple containing the AsyncioResult and the result of the task
        """
        impl = cls._trigger_or_join_task
        if retry_on_timeout:
            impl = tenacity.retry(
                stop=tenacity.stop_after_attempt(max_retries_on_timeout),
                retry=tenacity.retry_if_exception_type(
                    (
                        asyncio.TimeoutError,
                        TimeoutError,
                        TaskRevokedError,
                        MaxPendingInterruptsError,
                    )
                ),
                reraise=True,
            )(impl)
        return await impl(
            workflow,
            timeout=timeout,
            return_type=return_type,
            propagate=propagate,
            forget_cooldown=forget_cooldown,
            forget=forget,
            interrupt_every=interrupt_every,
            max_pending_interrupts=max_pending_interrupts,
            track_started=track_started,
        )
