from __future__ import annotations

"""Shared, thread-safe resource cache with in-flight coordination.

Example
-------
>>> cache = SharedResourceCache[str](maxsize=2)
>>> def build():
...     return "value"
>>> cache.get_or_create("key", build)
'value'
>>> cache.get_or_create("key", build)
'value'
>>>
>>> # Example: sharing a heavy ONNX Runtime session
>>> # from highlighter.predictors.ort.session import build_session  # pseudo-helper
>>> # session = cache.get_or_create(session_key, build_session)
"""

import threading
from collections import OrderedDict
from typing import Callable, Dict, Generic, Hashable, Optional, TypeVar

T = TypeVar("T")


class _Inflight(Generic[T]):
    __slots__ = ("event", "error", "value")

    def __init__(self) -> None:
        self.event = threading.Event()
        self.error: Optional[BaseException] = None
        self.value: Optional[T] = None


class SharedResourceCache(Generic[T]):
    """Thread-safe cache for shared heavyweight resources.

    - Ensures only one builder runs per key at a time.
    - Waiters block until the in-flight build completes.
    - LRU eviction enforces a maximum cache size (maxsize <= 0 disables eviction).
    """

    def __init__(self, maxsize: int = 128) -> None:
        self._cache: "OrderedDict[Hashable, T]" = OrderedDict()
        self._inflight: Dict[Hashable, _Inflight[T]] = {}
        self._lock = threading.RLock()
        self._maxsize = maxsize

    def _evict_if_needed(self, maxsize: int) -> None:
        if maxsize <= 0:
            return
        while len(self._cache) > maxsize:
            self._cache.popitem(last=False)

    def get_or_create(
        self,
        key: Hashable,
        create_fn: Callable[[], T],
        *,
        maxsize: Optional[int] = None,
    ) -> T:
        size = self._maxsize if maxsize is None else maxsize
        with self._lock:
            cached = self._cache.get(key)
            if cached is not None:
                self._cache.move_to_end(key)
                return cached

            inflight = self._inflight.get(key)
            if inflight is None:
                inflight = _Inflight[T]()
                self._inflight[key] = inflight
                creator = True
            else:
                creator = False

        if not creator:
            inflight.event.wait()
            if inflight.error:
                raise RuntimeError("Shared resource creation failed.") from inflight.error
            if inflight.value is None:
                raise RuntimeError("Shared cache creation finished without a value.")
            return inflight.value

        try:
            value = create_fn()
        except Exception as exc:
            with self._lock:
                inflight.error = exc
                self._inflight.pop(key, None)
                inflight.event.set()
            raise

        with self._lock:
            self._cache[key] = value
            self._evict_if_needed(int(size))

            inflight.value = value
            self._inflight.pop(key, None)
            inflight.event.set()

        return value

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

    def keys(self) -> list[Hashable]:
        with self._lock:
            return list(self._cache.keys())
