"""外部系统进程内运行态缓存与同步 single-flight。"""

from __future__ import annotations

import atexit
import time
from collections import OrderedDict
from concurrent.futures import Future
from contextlib import contextmanager
from dataclasses import dataclass
from threading import RLock
from typing import Any, Callable, Iterator, TypeVar


T = TypeVar("T")


@dataclass
class _RuntimeEntry:
    client: Any = None
    active_client_leases: int = 0
    evicted: bool = False
    auth_headers: dict[str, str] | None = None
    auth_expires_at: float = 0.0
    auth_generation: int = 0
    spec: dict[str, Any] | None = None
    spec_expires_at: float = 0.0
    catalog: Any = None
    catalog_spec: dict[str, Any] | None = None
    mcp_catalog: Any = None
    mcp_catalog_expires_at: float = 0.0


class ExternalRuntimeCache:
    """按连接身份隔离的有界 LRU；凭据、Token 和规格均只驻留当前进程。"""

    def __init__(self, *, max_entries: int = 256):
        self.max_entries = max_entries
        self._entries: OrderedDict[str, _RuntimeEntry] = OrderedDict()
        self._inflight: dict[tuple[str, str], Future[Any]] = {}
        self._lock = RLock()

    @staticmethod
    def _close_client(client: Any) -> None:
        close = getattr(client, "close", None)
        if callable(close):
            try:
                close()
            except Exception:
                pass

    def _entry_locked(self, identity: str) -> _RuntimeEntry:
        entry = self._entries.get(identity)
        if entry is None:
            entry = _RuntimeEntry()
            self._entries[identity] = entry
        self._entries.move_to_end(identity)
        while len(self._entries) > self.max_entries:
            _, evicted = self._entries.popitem(last=False)
            evicted.evicted = True
            if evicted.active_client_leases == 0:
                self._close_client(evicted.client)
        return entry

    @contextmanager
    def client(self, identity: str, factory: Callable[[], T]) -> Iterator[T]:
        with self._lock:
            entry = self._entry_locked(identity)
            if entry.client is None:
                entry.client = factory()
            entry.active_client_leases += 1
            client = entry.client
        try:
            yield client
        finally:
            close_client = None
            with self._lock:
                entry.active_client_leases -= 1
                if entry.evicted and entry.active_client_leases == 0:
                    close_client = entry.client
                    entry.client = None
            if close_client is not None:
                self._close_client(close_client)

    def get_auth(self, identity: str) -> dict[str, str] | None:
        now = time.monotonic()
        with self._lock:
            entry = self._entries.get(identity)
            if entry is None or entry.auth_expires_at <= now:
                if entry is not None:
                    entry.auth_headers = None
                    entry.auth_expires_at = 0.0
                return None
            self._entries.move_to_end(identity)
            return dict(entry.auth_headers or {})

    def set_auth(
        self, identity: str, headers: dict[str, str], *, ttl_seconds: float
    ) -> None:
        with self._lock:
            entry = self._entry_locked(identity)
            entry.auth_headers = dict(headers)
            entry.auth_expires_at = time.monotonic() + max(0.0, ttl_seconds)
            entry.auth_generation += 1

    def auth_state(self, identity: str) -> tuple[dict[str, str] | None, int]:
        headers = self.get_auth(identity)
        with self._lock:
            entry = self._entries.get(identity)
            return headers, entry.auth_generation if entry is not None else 0

    def invalidate_auth(self, identity: str) -> None:
        with self._lock:
            entry = self._entries.get(identity)
            if entry is not None:
                entry.auth_headers = None
                entry.auth_expires_at = 0.0

    def get_spec(self, identity: str) -> dict[str, Any] | None:
        now = time.monotonic()
        with self._lock:
            entry = self._entries.get(identity)
            if entry is None or entry.spec_expires_at <= now:
                if entry is not None:
                    entry.spec = None
                    entry.spec_expires_at = 0.0
                    entry.catalog = None
                    entry.catalog_spec = None
                return None
            self._entries.move_to_end(identity)
            return entry.spec

    def set_spec(
        self, identity: str, spec: dict[str, Any], *, ttl_seconds: float
    ) -> None:
        with self._lock:
            entry = self._entry_locked(identity)
            entry.spec = spec
            entry.spec_expires_at = time.monotonic() + max(0.0, ttl_seconds)
            entry.catalog = None
            entry.catalog_spec = None

    def invalidate_spec(self, identity: str) -> None:
        with self._lock:
            entry = self._entries.get(identity)
            if entry is not None:
                entry.spec = None
                entry.spec_expires_at = 0.0
                entry.catalog = None
                entry.catalog_spec = None

    def get_catalog(self, identity: str, spec: dict[str, Any]) -> Any:
        with self._lock:
            entry = self._entries.get(identity)
            if entry is None or entry.catalog_spec is not spec:
                return None
            self._entries.move_to_end(identity)
            return entry.catalog

    def set_catalog(self, identity: str, spec: dict[str, Any], catalog: Any) -> None:
        with self._lock:
            entry = self._entry_locked(identity)
            entry.catalog_spec = spec
            entry.catalog = catalog

    def get_mcp_catalog(self, identity: str) -> Any:
        now = time.monotonic()
        with self._lock:
            entry = self._entries.get(identity)
            if entry is None or entry.mcp_catalog_expires_at <= now:
                if entry is not None:
                    entry.mcp_catalog = None
                    entry.mcp_catalog_expires_at = 0.0
                return None
            self._entries.move_to_end(identity)
            return entry.mcp_catalog

    def set_mcp_catalog(
        self, identity: str, catalog: Any, *, ttl_seconds: float
    ) -> None:
        with self._lock:
            entry = self._entry_locked(identity)
            entry.mcp_catalog = catalog
            entry.mcp_catalog_expires_at = time.monotonic() + max(0.0, ttl_seconds)

    def invalidate_mcp_catalog(self, identity: str) -> None:
        with self._lock:
            entry = self._entries.get(identity)
            if entry is not None:
                entry.mcp_catalog = None
                entry.mcp_catalog_expires_at = 0.0

    def singleflight(self, namespace: str, key: str, compute: Callable[[], T]) -> T:
        flight_key = (namespace, key)
        with self._lock:
            future = self._inflight.get(flight_key)
            leader = future is None
            if leader:
                future = Future()
                self._inflight[flight_key] = future
        assert future is not None
        if not leader:
            return future.result()
        try:
            result = compute()
        except BaseException as exc:
            future.set_exception(exc)
            raise
        else:
            future.set_result(result)
            return result
        finally:
            with self._lock:
                if self._inflight.get(flight_key) is future:
                    self._inflight.pop(flight_key, None)

    def clear(self) -> None:
        with self._lock:
            entries = list(self._entries.values())
            self._entries.clear()
        for entry in entries:
            self._close_client(entry.client)

    def spec_count(self) -> int:
        now = time.monotonic()
        with self._lock:
            return sum(
                1
                for entry in self._entries.values()
                if entry.spec is not None and entry.spec_expires_at > now
            )


RUNTIME_CACHE = ExternalRuntimeCache()
atexit.register(RUNTIME_CACHE.clear)
