from __future__ import annotations

import asyncio
import threading
import time
from collections import OrderedDict
from collections.abc import Awaitable, Callable, Hashable, Iterator
from dataclasses import dataclass
from typing import Generic, TypeVar

K = TypeVar("K", bound=Hashable)
V = TypeVar("V")


def _now() -> float:
    return time.monotonic()


@dataclass(frozen=True, slots=True)
class _CacheItem(Generic[V]):
    value: V
    expires_at: float | None


class TTLCache(Generic[K, V]):
    """Thread-safe in-memory TTL cache with LRU eviction.

    Works correctly in both sync (threading) and async (asyncio) contexts:
    the internal ``RLock`` never blocks the event loop for any meaningful duration
    because all operations are O(1) / O(expired_keys).

    Does NOT solve "thundering herd" for async code — pair with
    ``AsyncSingleFlight`` for that.

    Args:
        ttl_seconds: ``None`` — entries live forever (until evicted by LRU or cleared).
            ``> 0`` — entries expire after that many seconds.
            ``0`` — cache is disabled: ``get`` always returns the default,
            ``set`` is a no-op.
        maxsize: ``None`` — unbounded (not recommended for long-running services).
            ``> 0`` — LRU eviction kicks in when the limit is reached.
        prune_every: Expired entries are eagerly pruned every *prune_every* write
            operations. Reads that hit an expired entry always remove it
            immediately regardless of this counter.
    """

    __slots__ = ("_data", "_lock", "_maxsize", "_prune_every", "_ttl", "_writes")

    def __init__(
        self,
        *,
        ttl_seconds: float | None = None,
        maxsize: int | None = 1024,
        prune_every: int = 128,
    ) -> None:
        if ttl_seconds is not None and ttl_seconds < 0:
            raise ValueError("ttl_seconds must be None or >= 0")
        if maxsize is not None and maxsize <= 0:
            raise ValueError("maxsize must be None or > 0")
        if prune_every <= 0:
            raise ValueError("prune_every must be > 0")

        self._ttl = ttl_seconds
        self._maxsize = maxsize
        self._data: OrderedDict[K, _CacheItem[V]] = OrderedDict()
        self._lock = threading.RLock()
        self._writes = 0
        self._prune_every = prune_every

    @property
    def enabled(self) -> bool:
        """``False`` when ttl_seconds == 0 (cache disabled)."""
        return self._ttl != 0

    @property
    def ttl_seconds(self) -> float | None:
        """Configured TTL in seconds, or ``None`` if entries never expire."""
        return self._ttl

    @property
    def maxsize(self) -> int | None:
        """Maximum number of entries before LRU eviction, or ``None`` if unbounded."""
        return self._maxsize

    def get(self, key: K, default: V | None = None) -> V | None:
        if self._ttl == 0:
            return default

        with self._lock:
            item = self._data.get(key)
            if item is None:
                return default
            if item.expires_at is not None and item.expires_at <= _now():
                del self._data[key]
                return default
            self._data.move_to_end(key, last=True)
            return item.value

    def set(self, key: K, value: V) -> None:
        if self._ttl == 0:
            return

        expires_at = None if self._ttl is None else (_now() + self._ttl)
        with self._lock:
            self._data[key] = _CacheItem(value=value, expires_at=expires_at)
            self._data.move_to_end(key, last=True)
            self._enforce_maxsize_locked()
            self._writes += 1
            if self._writes % self._prune_every == 0:
                self._prune_expired_locked()

    def pop(self, key: K, default: V | None = None) -> V | None:
        if self._ttl == 0:
            return default

        with self._lock:
            item = self._data.pop(key, None)
            if item is None:
                return default
            if item.expires_at is not None and item.expires_at <= _now():
                return default
            return item.value

    def delete(self, key: K) -> None:
        with self._lock:
            self._data.pop(key, None)

    def clear(self) -> None:
        with self._lock:
            self._data.clear()
            self._writes = 0

    def __contains__(self, key: object) -> bool:
        if self._ttl == 0:
            return False
        with self._lock:
            item = self._data.get(key)  # type: ignore[arg-type]
            if item is None:
                return False
            if item.expires_at is not None and item.expires_at <= _now():
                del self._data[key]  # type: ignore[arg-type]
                return False
            return True

    def __len__(self) -> int:
        if self._ttl == 0:
            return 0
        with self._lock:
            self._prune_expired_locked()
            return len(self._data)

    def __iter__(self) -> Iterator[K]:
        """Iterate over *live* keys (snapshot, safe to use outside the lock)."""
        with self._lock:
            self._prune_expired_locked()
            return iter(list(self._data.keys()))

    def __repr__(self) -> str:
        return (
            f"{type(self).__name__}("
            f"ttl_seconds={self._ttl!r}, "
            f"maxsize={self._maxsize!r}, "
            f"size={len(self._data)}"
            f")"
        )

    def stats(self) -> dict[str, int | None]:
        with self._lock:
            self._prune_expired_locked()
            return {
                "size": len(self._data),
                "maxsize": self._maxsize,
                "writes": self._writes,
            }

    def _enforce_maxsize_locked(self) -> None:
        if self._maxsize is None:
            return
        while len(self._data) > self._maxsize:
            self._data.popitem(last=False)

    def _prune_expired_locked(self) -> None:
        if not self._data:
            return
        now = _now()
        to_delete = [
            k
            for k, item in self._data.items()
            if item.expires_at is not None and item.expires_at <= now
        ]
        for k in to_delete:
            del self._data[k]


class AsyncSingleFlight(Generic[K]):
    """Thundering-herd guard for async code.

    When multiple coroutines request the same key simultaneously, only the
    first one does the actual work; the rest await the same lock and then
    re-check the cache themselves (double-checked locking pattern).

    This class does **not** cache results — combine it with ``TTLCache``.

    Args:
        max_locks: Upper bound on the number of in-flight per-key locks kept in
            memory. When the limit is reached, idle locks are evicted (FIFO) to
            make room. Active locks are never evicted.

    Example:
        >>> async with singleflight.lock(key):
        ...     cached = cache.get(key)
        ...     if cached is not None:
        ...         return cached
        ...     result = await fetch_from_api(key)
        ...     cache.set(key, result)
        ...     return result
    """

    __slots__ = ("_locks", "_max_locks", "_meta_lock")

    def __init__(self, *, max_locks: int = 4096) -> None:
        if max_locks <= 0:
            raise ValueError("max_locks must be > 0")
        self._locks: OrderedDict[K, asyncio.Lock] = OrderedDict()
        self._meta_lock = asyncio.Lock()
        self._max_locks = max_locks

    def __repr__(self) -> str:
        return (
            f"{type(self).__name__}("
            f"in_flight={len(self._locks)}, "
            f"max_locks={self._max_locks}"
            f")"
        )

    class _KeyLockCtx:
        __slots__ = ("_acquired_lock", "_key", "_sf")

        def __init__(self, sf: AsyncSingleFlight[K], key: K) -> None:
            self._sf = sf
            self._key = key
            self._acquired_lock: asyncio.Lock | None = None

        async def __aenter__(self) -> None:
            self._acquired_lock = await self._sf._acquire(self._key)

        async def __aexit__(
            self,
            exc_type: type[BaseException] | None,
            exc: BaseException | None,
            tb: object,
        ) -> None:
            assert self._acquired_lock is not None
            self._acquired_lock.release()
            await self._sf._cleanup(self._key, self._acquired_lock)

    def lock(self, key: K) -> AsyncSingleFlight._KeyLockCtx:
        return AsyncSingleFlight._KeyLockCtx(self, key)

    async def _acquire(self, key: K) -> asyncio.Lock:
        async with self._meta_lock:
            lock_obj = self._locks.get(key)
            if lock_obj is None:
                await self._maybe_evict_locked()
                lock_obj = asyncio.Lock()
                self._locks[key] = lock_obj
            else:
                self._locks.move_to_end(key, last=True)
        await lock_obj.acquire()
        return lock_obj

    async def _cleanup(self, key: K, lock_obj: asyncio.Lock) -> None:
        """Remove the key from the registry if *our* lock is still there and idle."""
        async with self._meta_lock:
            cur = self._locks.get(key)
            if cur is not None and cur is lock_obj and not cur.locked():
                del self._locks[key]

    async def _maybe_evict_locked(self) -> None:
        """Evict oldest *idle* locks until under max_locks.  Active locks are skipped."""
        if len(self._locks) < self._max_locks:
            return
        evicted = 0
        for k, lk in list(self._locks.items()):
            if not lk.locked():
                del self._locks[k]
                evicted += 1
                if len(self._locks) < self._max_locks:
                    break


class SyncKeyedLock(Generic[K]):
    """Per-key locking for threaded code — sync analogue of ``AsyncSingleFlight``.

    Prevents multiple threads from concurrently executing the same expensive
    operation (e.g. a cache-miss network request) for the same key.

    Args:
        max_locks: Upper bound on idle per-key locks kept in memory.
            Idle locks are evicted FIFO when the limit is reached.
            Locks held by any thread are never evicted.
    """

    __slots__ = ("_locks", "_max_locks", "_meta_lock")

    def __init__(self, *, max_locks: int = 4096) -> None:
        if max_locks <= 0:
            raise ValueError("max_locks must be > 0")
        self._locks: OrderedDict[K, threading.Lock] = OrderedDict()
        self._meta_lock = threading.RLock()
        self._max_locks = max_locks

    def __repr__(self) -> str:
        return (
            f"{type(self).__name__}("
            f"tracked={len(self._locks)}, "
            f"max_locks={self._max_locks}"
            f")"
        )

    class _Ctx:
        __slots__ = ("_key", "_lk", "_parent")

        def __init__(self, parent: SyncKeyedLock[K], key: K) -> None:
            self._parent = parent
            self._key = key
            self._lk: threading.Lock | None = None

        def __enter__(self) -> None:
            self._lk = self._parent._acquire(self._key)

        def __exit__(
            self,
            exc_type: type[BaseException] | None,
            exc: BaseException | None,
            tb: object,
        ) -> None:
            assert self._lk is not None
            self._lk.release()
            self._parent._cleanup(self._key, self._lk)

    def lock(self, key: K) -> SyncKeyedLock._Ctx:
        return SyncKeyedLock._Ctx(self, key)

    def _acquire(self, key: K) -> threading.Lock:
        with self._meta_lock:
            lk = self._locks.get(key)
            if lk is None:
                self._maybe_evict_locked()
                lk = threading.Lock()
                self._locks[key] = lk
            else:
                self._locks.move_to_end(key, last=True)
        lk.acquire()
        return lk

    def _cleanup(self, key: K, lk: threading.Lock) -> None:
        """Remove the key from the registry if our lock is still there and idle."""
        with self._meta_lock:
            cur = self._locks.get(key)
            if cur is not None and cur is lk and not cur.locked():
                del self._locks[key]

    def _maybe_evict_locked(self) -> None:
        """Evict oldest *idle* locks until under max_locks.  Active locks are skipped."""
        if len(self._locks) < self._max_locks:
            return
        for k, lk in list(self._locks.items()):
            if not lk.locked():
                del self._locks[k]
                if len(self._locks) < self._max_locks:
                    break


class AsyncResolver(Generic[K, V]):
    """Async TTL cache with thundering-herd protection.

    Wraps a fetcher callable with ``TTLCache`` and ``AsyncSingleFlight`` so that
    concurrent requests for the same key trigger only one underlying call.

    Args:
        fetcher: Async callable that produces a value for a given key.
        ttl_seconds: Cache entry lifetime; ``None`` means no expiry.
        maxsize: Maximum number of cached entries (LRU eviction).
    """

    __slots__ = ("_cache", "_fetcher", "_sf")

    def __init__(
        self,
        fetcher: Callable[[K], Awaitable[V]],
        *,
        ttl_seconds: float | None = None,
        maxsize: int | None = 256,
    ) -> None:
        self._cache: TTLCache[K, V] = TTLCache(ttl_seconds=ttl_seconds, maxsize=maxsize)
        self._sf: AsyncSingleFlight[K] = AsyncSingleFlight()
        self._fetcher = fetcher

    async def get(self, key: K) -> V:
        cached = self._cache.get(key)
        if cached is not None:
            return cached

        async with self._sf.lock(key):
            cached = self._cache.get(key)
            if cached is not None:
                return cached

            value = await self._fetcher(key)
            self._cache.set(key, value)
            return value

    def invalidate(self, key: K) -> None:
        self._cache.delete(key)

    def clear(self) -> None:
        self._cache.clear()


class SyncResolver(Generic[K, V]):
    """Sync TTL cache with per-key locking.

    Wraps a fetcher callable with ``TTLCache`` and ``SyncKeyedLock`` so that
    concurrent threads requesting the same key trigger only one underlying call.

    Args:
        fetcher: Callable that produces a value for a given key.
        ttl_seconds: Cache entry lifetime; ``None`` means no expiry.
        maxsize: Maximum number of cached entries (LRU eviction).
    """

    __slots__ = ("_cache", "_fetcher", "_lock")

    def __init__(
        self,
        fetcher: Callable[[K], V],
        *,
        ttl_seconds: float | None = None,
        maxsize: int | None = 256,
    ) -> None:
        self._cache: TTLCache[K, V] = TTLCache(ttl_seconds=ttl_seconds, maxsize=maxsize)
        self._lock: SyncKeyedLock[K] = SyncKeyedLock()
        self._fetcher = fetcher

    def get(self, key: K) -> V:
        cached = self._cache.get(key)
        if cached is not None:
            return cached

        with self._lock.lock(key):
            cached = self._cache.get(key)
            if cached is not None:
                return cached

            value = self._fetcher(key)
            self._cache.set(key, value)
            return value

    def invalidate(self, key: K) -> None:
        self._cache.delete(key)

    def clear(self) -> None:
        self._cache.clear()


__all__ = ["AsyncSingleFlight", "SyncKeyedLock", "TTLCache"]
