from collections.abc import Awaitable, Callable
from functools import wraps
from typing import Concatenate, ParamSpec, Protocol, TypeVar, cast

from audience_common.cache.backends import AsyncCache
from audience_common.cache.utils import make_key


class HasCacheProtocol(Protocol):
    cache: AsyncCache


HasCache = TypeVar("HasCache", bound=HasCacheProtocol)
P = ParamSpec("P")
R = TypeVar("R")


def cached(
    key: str,
    *,
    timeout: int | Callable[[], int] | None = None,
    timeout_param: str | None = None,
) -> Callable[
    [Callable[Concatenate[HasCache, P], Awaitable[R]]],
    Callable[Concatenate[HasCache, P], Awaitable[R]],
]:
    def decorator(
        func: Callable[Concatenate[HasCache, P], Awaitable[R]]
    ) -> Callable[Concatenate[HasCache, P], Awaitable[R]]:
        @wraps(func)
        async def wrapper(self: HasCache, /, *args: P.args, **kwargs: P.kwargs) -> R:
            cache_timeout = timeout
            if cache_timeout is None and timeout_param:
                cache_timeout = getattr(self, timeout_param, None)
            cache_key = make_key(key, *args, **kwargs)
            value = await self.cache.get(cache_key)
            if value is None:
                value = await func(self, *args, **kwargs)
                if callable(cache_timeout):
                    cache_timeout = cache_timeout()
                await self.cache.set(cache_key, value, timeout=cache_timeout)
            return cast(R, value)

        return wrapper

    return decorator