"""
RBSP Query Cache Layer.

Provides LRU caching with time-to-live expiry for query results, avoiding
redundant scoring of identical queries.  Two variants are included:

    - ``LRUCache``      Thread-safe LRU cache with configurable max size and TTL.
    - ``SegmentCache``  Extends LRU with generation-based invalidation so that
                        stale results are discarded when the underlying index
                        changes.

Inspired by Elasticsearch's per-segment query caching strategy: cached results
are only valid for the index generation that produced them.
"""

from __future__ import annotations

import threading
import time
from collections import OrderedDict
from typing import Any, Optional


class LRUCache:
    """Thread-safe LRU cache with time-to-live expiry.

    Parameters
    ----------
    max_size:
        Maximum number of entries before the least-recently-used item is
        evicted.  Must be at least 1.
    ttl:
        Time-to-live in seconds for each entry.  Entries older than *ttl*
        are treated as missing on ``get`` and lazily removed.
    """

    __slots__ = (
        "_max_size",
        "_ttl",
        "_data",
        "_lock",
        "_hit_count",
        "_miss_count",
        "_eviction_count",
    )

    def __init__(self, max_size: int = 1024, ttl: float = 300.0) -> None:
        if max_size < 1:
            raise ValueError("max_size must be >= 1")
        if ttl <= 0.0:
            raise ValueError("ttl must be > 0")
        self._max_size = max_size
        self._ttl = ttl
        # OrderedDict: key -> (value, expiry_time)
        self._data: OrderedDict[str, tuple[Any, float]] = OrderedDict()
        self._lock = threading.Lock()
        self._hit_count = 0
        self._miss_count = 0
        self._eviction_count = 0

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    def get(self, key: str) -> Optional[Any]:
        """Return the cached value for *key*, or ``None`` if missing/expired."""
        with self._lock:
            entry = self._data.get(key)
            if entry is None:
                self._miss_count += 1
                return None
            value, expiry = entry
            if time.monotonic() > expiry:
                # Expired — remove lazily.
                del self._data[key]
                self._miss_count += 1
                return None
            # Move to most-recently-used end.
            self._data.move_to_end(key)
            self._hit_count += 1
            return value

    def put(self, key: str, value: Any) -> None:
        """Insert or update *key* with *value*, evicting LRU if necessary."""
        expiry = time.monotonic() + self._ttl
        with self._lock:
            if key in self._data:
                # Update existing — move to end.
                self._data.move_to_end(key)
                self._data[key] = (value, expiry)
            else:
                self._data[key] = (value, expiry)
                if len(self._data) > self._max_size:
                    self._data.popitem(last=False)
                    self._eviction_count += 1

    def invalidate(self, key: str) -> None:
        """Remove a specific entry.  No-op if *key* is absent."""
        with self._lock:
            self._data.pop(key, None)

    def clear(self) -> None:
        """Remove all entries and reset stats."""
        with self._lock:
            self._data.clear()
            self._hit_count = 0
            self._miss_count = 0
            self._eviction_count = 0

    def stats(self) -> dict:
        """Return cache performance statistics."""
        with self._lock:
            total = self._hit_count + self._miss_count
            return {
                "hit_count": self._hit_count,
                "miss_count": self._miss_count,
                "eviction_count": self._eviction_count,
                "size": len(self._data),
                "hit_rate": self._hit_count / total if total > 0 else 0.0,
            }

    # ------------------------------------------------------------------
    # Dunder helpers
    # ------------------------------------------------------------------

    def __contains__(self, key: object) -> bool:
        with self._lock:
            entry = self._data.get(key)  # type: ignore[arg-type]
            if entry is None:
                return False
            _, expiry = entry
            return time.monotonic() <= expiry

    def __len__(self) -> int:
        with self._lock:
            return len(self._data)


class SegmentCache(LRUCache):
    """LRU cache with generation-based invalidation.

    When the underlying index mutates (documents added/removed), the caller
    bumps the generation counter.  Cached results tagged with an older
    generation are treated as stale and discarded on access.

    Parameters
    ----------
    max_size:
        Maximum number of cached entries.
    ttl:
        Per-entry time-to-live in seconds.
    """

    __slots__ = ("_generation",)

    def __init__(self, max_size: int = 1024, ttl: float = 300.0) -> None:
        super().__init__(max_size=max_size, ttl=ttl)
        self._generation = 0

    # ------------------------------------------------------------------
    # Generation management
    # ------------------------------------------------------------------

    def set_generation(self, generation: int) -> None:
        """Advance the index generation.

        Entries cached under older generations will be lazily invalidated
        on the next ``get``.
        """
        with self._lock:
            self._generation = generation

    # ------------------------------------------------------------------
    # Overridden API
    # ------------------------------------------------------------------

    def put(self, key: str, value: Any, generation: int | None = None) -> None:  # type: ignore[override]
        """Cache *value* under *key* tagged with *generation*.

        If *generation* is ``None`` the current generation is used.
        """
        if generation is None:
            generation = self._generation
        expiry = time.monotonic() + self._ttl
        with self._lock:
            if key in self._data:
                self._data.move_to_end(key)
                self._data[key] = (value, expiry, generation)
            else:
                self._data[key] = (value, expiry, generation)
                if len(self._data) > self._max_size:
                    self._data.popitem(last=False)
                    self._eviction_count += 1

    def get(self, key: str, current_generation: int | None = None) -> Optional[Any]:  # type: ignore[override]
        """Return cached value only if generation matches and TTL is valid.

        If *current_generation* is ``None`` the cache's internal generation
        counter is used.
        """
        if current_generation is None:
            current_generation = self._generation
        with self._lock:
            entry = self._data.get(key)
            if entry is None:
                self._miss_count += 1
                return None
            # Entries are 3-tuples: (value, expiry, generation).
            value, expiry, gen = entry
            if gen != current_generation or time.monotonic() > expiry:
                del self._data[key]
                self._miss_count += 1
                return None
            self._data.move_to_end(key)
            self._hit_count += 1
            return value
