import threading
import typing
from collections.abc import Iterable, Iterator, MutableMapping

from .misc import SENTINEL

__all__ = ['LRU']

K = typing.TypeVar('K')
V = typing.TypeVar('V')


class LRU(MutableMapping[K, V], typing.Generic[K, V]):
    """
    Implementation of a length-limited LRU map.

    The mapping is thread-safe, and internally uses a lock to avoid concurrency
    issues. However, access operations like ``lru[key]`` are fast and
    lock-free.
    """

    __slots__ = ('_count', '_lock', '_ordering', '_values')

    def __init__(self, count: int, pairs: Iterable[tuple[K, V]] = ()):
        assert count > 0, "LRU needs a positive count"
        self._count = count
        self._lock = threading.RLock()
        self._values: dict[K, V] = {}
        #
        # The dict self._values contains the LRU items, while self._ordering
        # only keeps track of their order, the most recently used ones being
        # last. For performance reasons, we only use the lock when modifying
        # the LRU, while reading it is lock-free (and thus faster).
        #
        # This strategy may result in inconsistencies between self._values and
        # self._ordering. Indeed, concurrently accessed keys may be missing
        # from self._ordering, but will eventually be added. This could result
        # in keys being added back in self._ordering after their actual removal
        # from the LRU. This results in the following invariant:
        #
        #     self._values <= self._ordering | "keys being accessed"
        #
        self._ordering: dict[K, None] = {}

        # Initialize
        for key, value in pairs:
            self[key] = value

    @property
    def count(self) -> int:
        return self._count

    def __contains__(self, key: object) -> bool:
        return key in self._values

    def __getitem__(self, key: K) -> V:
        val = self._values[key]
        # move key at the last position in self._ordering
        self._ordering[key] = self._ordering.pop(key, None)
        return val

    def __setitem__(self, key: K, val: V):
        values = self._values
        ordering = self._ordering
        with self._lock:
            values[key] = val
            ordering[key] = ordering.pop(key, None)
            while True:
                # if we have too many keys in ordering, filter them out
                if len(ordering) > len(values):
                    # (copy to avoid concurrent changes on ordering)
                    for k in ordering.copy():
                        if k not in values:
                            ordering.pop(k, None)
                # check if we have too many keys
                if len(values) <= self._count:
                    break
                # if so, pop the least recently used
                try:
                    # have a default in case of concurrent accesses
                    key = next(iter(ordering), key)
                except RuntimeError:
                    # ordering modified during iteration, retry
                    continue
                values.pop(key, None)
                ordering.pop(key, None)

    def __delitem__(self, key: K):
        self.pop(key)

    def __len__(self) -> int:
        return len(self._values)

    def __iter__(self) -> Iterator[K]:
        return iter(self.snapshot)

    @property
    def snapshot(self) -> dict[K, V]:
        """ Return a copy of the LRU (ordered according to LRU first). """
        with self._lock:
            values = self._values
            # build result in expected order (copy self._ordering to avoid concurrent changes)
            result = {
                key: val
                for key in self._ordering.copy()
                if (val := values.get(key, SENTINEL)) is not SENTINEL
            }
            if len(result) < len(values):
                # keys in value were missing from self._ordering, add them
                result.update(values)
        return result

    def pop(self, key: K, /, default=SENTINEL) -> V:
        with self._lock:
            self._ordering.pop(key, None)
            if default is SENTINEL:
                return self._values.pop(key)
            return self._values.pop(key, default)

    def clear(self):
        with self._lock:
            self._ordering.clear()
            self._values.clear()