"""Asynchronous S3-FIFO cache decorator, modeled after async_lru.alru_cache."""

from __future__ import annotations

import asyncio
import dataclasses
import inspect
import sys
from collections.abc import Callable, Coroutine, Hashable
from functools import _CacheInfo, _make_key, partial, partialmethod, update_wrapper
from typing import Any, Generic, TypedDict, TypeVar, Union, cast, final, overload

from s3fifo._s3fifo import GHOST_SIZE_RATIO, MOVE_TO_MAIN_THRESHOLD, SMALL_SIZE_RATIO, S3FIFOCache

if sys.version_info >= (3, 11):
    from typing import Self
else:
    Self = Any

if sys.version_info < (3, 14):
    from asyncio.coroutines import _is_coroutine  # type: ignore[attr-defined]

_T = TypeVar("_T")
_R = TypeVar("_R")
_Coro = Coroutine[Any, Any, _R]
_CB = Callable[..., _Coro[_R]]
_CBP = Union[_CB[_R], "partial[_Coro[_R]]", "partialmethod[_Coro[_R]]"]


@final
class _CacheParameters(TypedDict):
    maxsize: int
    typed: bool
    small_size_ratio: float
    ghost_size_ratio: float
    move_to_main_threshold: int
    tasks: int
    closed: bool


class _CacheClosedError(RuntimeError):
    def __init__(self, wrapper: object) -> None:
        super().__init__(f"as3fifo_cache is closed for {wrapper}")


class _EventLoopMismatchError(RuntimeError):
    def __init__(self) -> None:
        super().__init__(
            "as3fifo_cache is not safe to use across event loops: this cache "
            "instance was first used with a different event loop. "
            "Use separate cache instances per event loop."
        )


class _CoroutineFunctionRequiredError(RuntimeError):
    def __init__(self, fn: Any) -> None:
        super().__init__(f"Coroutine function is required, got {fn!r}")


class _InvalidMaxsizeError(ValueError):
    def __init__(self) -> None:
        super().__init__("maxsize must be >= 1")


class _UnsupportedDecoratorTargetError(NotImplementedError):
    def __init__(self, fn: Any) -> None:
        super().__init__(f"{fn!r} decorating is not supported")


@final
@dataclasses.dataclass(slots=True)
class _CacheItem(Generic[_R]):
    task: asyncio.Task[_R]
    waiters: int


@final
class _S3FIFOCacheWrapper(Generic[_R]):
    def __init__(
        self,
        fn: _CB[_R],
        maxsize: int,
        typed: bool,
        *,
        small_size_ratio: float,
        ghost_size_ratio: float,
        move_to_main_threshold: int,
    ) -> None:
        update_wrapper(self, fn)
        if sys.version_info < (3, 14):
            self._is_coroutine = _is_coroutine
        self.__wrapped__ = fn
        self.__maxsize = maxsize
        self.__typed = typed
        self.__small_size_ratio = small_size_ratio
        self.__ghost_size_ratio = ghost_size_ratio
        self.__move_to_main_threshold = move_to_main_threshold
        self.__cache = S3FIFOCache[_CacheItem[_R]](
            maxsize,
            small_size_ratio=small_size_ratio,
            ghost_size_ratio=ghost_size_ratio,
            move_to_main_threshold=move_to_main_threshold,
        )
        self.__closed = False
        self.__hits = 0
        self.__misses = 0
        self.__first_loop: asyncio.AbstractEventLoop | None = None

    @property
    def __tasks(self) -> list[asyncio.Task[_R]]:
        return list({cache_item.task for cache_item in self.__cache.values() if not cache_item.task.done()})

    def _check_loop(self, loop: asyncio.AbstractEventLoop) -> None:
        if self.__first_loop is None:
            self.__first_loop = loop
        elif self.__first_loop is not loop:
            raise _EventLoopMismatchError()

    def cache_invalidate(self, /, *args: Hashable, **kwargs: Any) -> bool:
        key = _make_key(args, kwargs, self.__typed)
        return self.__cache.remove(key, include_ghost=False)

    def cache_clear(self) -> None:
        self.__hits = 0
        self.__misses = 0
        self.__cache.clear()

    async def cache_close(self, *, wait: bool = False) -> None:
        loop = asyncio.get_running_loop()
        self._check_loop(loop)
        self.__closed = True

        tasks = self.__tasks
        if not tasks:
            return

        if not wait:
            for task in tasks:
                if not task.done():
                    task.cancel()

        await asyncio.gather(*tasks, return_exceptions=True)

    def cache_info(self) -> _CacheInfo:
        return _CacheInfo(
            self.__hits,
            self.__misses,
            self.__maxsize,
            len(self.__cache),
        )

    def cache_parameters(self) -> _CacheParameters:
        return _CacheParameters(
            maxsize=self.__maxsize,
            typed=self.__typed,
            small_size_ratio=self.__small_size_ratio,
            ghost_size_ratio=self.__ghost_size_ratio,
            move_to_main_threshold=self.__move_to_main_threshold,
            tasks=len(self.__tasks),
            closed=self.__closed,
        )

    def _task_done_callback(self, key: Hashable, task: asyncio.Task[_R]) -> None:
        if task.cancelled() or getattr(task, "_exception", None) is not None:
            self.__cache.remove(key)

    async def _shield_and_handle_cancelled_error(self, cache_item: _CacheItem[_R], key: Hashable) -> _R:
        task = cache_item.task
        try:
            return await asyncio.shield(task)
        except asyncio.CancelledError:
            if cache_item.waiters == 1 and not task.done():
                task.cancel()
                self.__cache.remove(key)
            raise
        finally:
            cache_item.waiters -= 1

    async def __call__(self, /, *fn_args: Any, **fn_kwargs: Any) -> _R:
        if self.__closed:
            raise _CacheClosedError(self)

        loop = asyncio.get_running_loop()
        self._check_loop(loop)
        key = _make_key(fn_args, fn_kwargs, self.__typed)

        cache_item = self.__cache.get(key)
        if cache_item is not None:
            self.__hits += 1
            if not cache_item.task.done():
                cache_item.waiters += 1
                return await self._shield_and_handle_cancelled_error(cache_item, key)
            return cache_item.task.result()

        coro = self.__wrapped__(*fn_args, **fn_kwargs)
        task: asyncio.Task[_R] = loop.create_task(coro)
        task.add_done_callback(partial(self._task_done_callback, key))

        cache_item = _CacheItem(task=task, waiters=1)
        self.__cache.put(key, cache_item)

        self.__misses += 1

        return await self._shield_and_handle_cancelled_error(cache_item, key)

    def __get__(self, instance: _T, owner: type[_T] | None) -> Self | _S3FIFOCacheWrapperInstanceMethod[_R, _T]:
        if owner is None:
            return self
        return _S3FIFOCacheWrapperInstanceMethod(self, instance)


@final
class _S3FIFOCacheWrapperInstanceMethod(Generic[_R, _T]):
    def __init__(
        self,
        wrapper: _S3FIFOCacheWrapper[_R],
        instance: _T,
    ) -> None:
        update_wrapper(self, wrapper)
        if sys.version_info < (3, 14):
            self._is_coroutine = _is_coroutine
        self.__wrapped__ = wrapper.__wrapped__
        self.__instance = instance
        self.__wrapper = wrapper

    def cache_invalidate(self, /, *args: Hashable, **kwargs: Any) -> bool:
        return self.__wrapper.cache_invalidate(self.__instance, *args, **kwargs)

    def cache_clear(self) -> None:
        self.__wrapper.cache_clear()

    async def cache_close(self, *, wait: bool = False) -> None:
        await self.__wrapper.cache_close(wait=wait)

    def cache_info(self) -> _CacheInfo:
        return self.__wrapper.cache_info()

    def cache_parameters(self) -> _CacheParameters:
        return self.__wrapper.cache_parameters()

    async def __call__(self, /, *fn_args: Any, **fn_kwargs: Any) -> _R:
        return await self.__wrapper(self.__instance, *fn_args, **fn_kwargs)


def _make_async_wrapper(
    maxsize: int,
    typed: bool,
    *,
    small_size_ratio: float,
    ghost_size_ratio: float,
    move_to_main_threshold: int,
) -> Callable[[_CBP[_R]], _S3FIFOCacheWrapper[_R]]:
    def wrapper(fn: _CBP[_R]) -> _S3FIFOCacheWrapper[_R]:
        origin = fn

        while isinstance(origin, (partial, partialmethod)):
            origin = origin.func

        if not inspect.iscoroutinefunction(origin):
            raise _CoroutineFunctionRequiredError(fn)

        if hasattr(fn, "_make_unbound_method"):
            fn = fn._make_unbound_method()  # type: ignore[union-attr]

        wrapped = _S3FIFOCacheWrapper(
            cast(_CB[_R], fn),
            maxsize,
            typed,
            small_size_ratio=small_size_ratio,
            ghost_size_ratio=ghost_size_ratio,
            move_to_main_threshold=move_to_main_threshold,
        )
        if sys.version_info >= (3, 12):
            wrapped = inspect.markcoroutinefunction(wrapped)
        return wrapped

    return wrapper


@overload
def as3fifo_cache(
    maxsize: int = 128,
    typed: bool = False,
    *,
    small_size_ratio: float = SMALL_SIZE_RATIO,
    ghost_size_ratio: float = GHOST_SIZE_RATIO,
    move_to_main_threshold: int = MOVE_TO_MAIN_THRESHOLD,
) -> Callable[[_CBP[_R]], _S3FIFOCacheWrapper[_R]]: ...


@overload
def as3fifo_cache(
    maxsize: _CBP[_R],
    /,
    typed: bool = False,
    *,
    small_size_ratio: float = SMALL_SIZE_RATIO,
    ghost_size_ratio: float = GHOST_SIZE_RATIO,
    move_to_main_threshold: int = MOVE_TO_MAIN_THRESHOLD,
) -> _S3FIFOCacheWrapper[_R]: ...


def as3fifo_cache(
    maxsize: int | _CBP[_R] = 128,
    typed: bool = False,
    *,
    small_size_ratio: float = SMALL_SIZE_RATIO,
    ghost_size_ratio: float = GHOST_SIZE_RATIO,
    move_to_main_threshold: int = MOVE_TO_MAIN_THRESHOLD,
) -> Callable[[_CBP[_R]], _S3FIFOCacheWrapper[_R]] | _S3FIFOCacheWrapper[_R]:
    """S3-FIFO cache decorator for async functions.

    Drop-in replacement for async_lru.alru_cache using the S3-FIFO eviction
    algorithm instead of LRU.

    Usage::

        @as3fifo_cache
        async def func(x):
            ...

        @as3fifo_cache(maxsize=256)
        async def func(x):
            ...
    """
    if isinstance(maxsize, int):
        if maxsize < 1:
            raise _InvalidMaxsizeError()
        return _make_async_wrapper(
            maxsize,
            typed,
            small_size_ratio=small_size_ratio,
            ghost_size_ratio=ghost_size_ratio,
            move_to_main_threshold=move_to_main_threshold,
        )

    fn = cast(_CB[_R], maxsize)
    if callable(fn) or hasattr(fn, "_make_unbound_method"):
        return _make_async_wrapper(
            128,
            typed,
            small_size_ratio=small_size_ratio,
            ghost_size_ratio=ghost_size_ratio,
            move_to_main_threshold=move_to_main_threshold,
        )(fn)
    raise _UnsupportedDecoratorTargetError(fn)
