"""
cortex/core/working_memory.py — Working Memory Cache (Issue #371).

Fast in-memory session store for current session context. Provides sub-millisecond
access to session-local facts, task state, and recent turn summaries — checked
BEFORE any DB query in the retrieval pipeline.

Design decisions:
    - Session-scoped: all entries are keyed by (session_id, key) so multiple
      concurrent sessions share one cache instance safely.
    - Thread-safe: uses asyncio.Lock; all public methods are coroutines.
    - LRU eviction: per-session doubly-linked list; O(1) get/put operations.
    - TTL: checked lazily on access (no background cleanup thread required).
    - Size cap: per-session item count, not global, so sessions don't starve
      each other.  Callers should invoke clear_session() when a session ends
      to reclaim memory.

Categories (matches Issue #371 spec):
    task_state          Active task / current step
    project_state       Current project facts
    session_facts       Session-local facts (not yet in DB)
    turn_summaries      Per-turn summaries injected into context
    entities            Active entities / people / projects
    open_loops          Outstanding commitments / questions
    tool_state          Current tool invocation state
    injected_bundles    Bundles already injected this session (dedup guard)

Dependencies: None (pure stdlib — no cortex imports to avoid circular deps).
Dependents: cortex.core.retrieval (wired in at retrieve() entry point).
"""

from __future__ import annotations

import asyncio
import logging
import time
from collections import OrderedDict
from dataclasses import dataclass, field
from enum import Enum
from typing import Any, Dict, List, Optional

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Category enum
# ---------------------------------------------------------------------------

class WorkingMemoryCategory(str, Enum):
    """Valid categories for working memory entries.

    Maps directly to the eight categories defined in Issue #371.
    Using a str Enum means values can be compared to plain strings:
        entry.category == "task_state"  → True
    """

    TASK_STATE = "task_state"
    PROJECT_STATE = "project_state"
    SESSION_FACTS = "session_facts"
    TURN_SUMMARIES = "turn_summaries"
    ENTITIES = "entities"
    OPEN_LOOPS = "open_loops"
    TOOL_STATE = "tool_state"
    INJECTED_BUNDLES = "injected_bundles"

    @classmethod
    def values(cls) -> frozenset[str]:
        """Return all category string values as a frozen set."""
        return frozenset(m.value for m in cls)


# ---------------------------------------------------------------------------
# Internal entry type
# ---------------------------------------------------------------------------

@dataclass
class _CacheEntry:
    """A single item stored in the working memory cache.

    Attributes:
        key: Caller-supplied string key (unique within a session).
        value: Arbitrary payload. Kept as-is; callers are responsible for
               serialisability if they need to persist entries later.
        category: One of the eight WorkingMemoryCategory values.
        session_id: Owning session — used for session-scoped queries and clear.
        created_at: Monotonic timestamp (float seconds) when the entry was put.
        expires_at: Monotonic timestamp after which the entry is expired.
                    None means no TTL (session-end only eviction).
        last_accessed: Monotonic timestamp of last successful get() — drives LRU.
    """
    key: str
    value: Any
    category: str
    session_id: str
    created_at: float
    expires_at: Optional[float]
    last_accessed: float


# ---------------------------------------------------------------------------
# LRU container (per-session)
# ---------------------------------------------------------------------------

class _SessionLRU:
    """OrderedDict-based LRU store for a single session.

    Keys are the caller's ``key`` strings. move_to_end(key) promotes
    accessed items; popitem(last=False) evicts the LRU entry.
    """

    def __init__(self, max_items: int) -> None:
        self._max = max_items
        self._data: OrderedDict[str, _CacheEntry] = OrderedDict()

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

    def get(self, key: str) -> Optional[_CacheEntry]:
        """Retrieve entry by key, promoting it to most-recently-used."""
        entry = self._data.get(key)
        if entry is not None:
            self._data.move_to_end(key)  # mark as recently used
        return entry

    def put(self, entry: _CacheEntry) -> Optional[_CacheEntry]:
        """Insert or update entry. Returns the evicted entry (if any) or None."""
        # Guard: if max capacity < 1, refuse to store (audit fix: max_items=0 crash)
        if self._max < 1:
            return None

        if entry.key in self._data:
            self._data.move_to_end(entry.key)
            self._data[entry.key] = entry
            return None

        evicted: Optional[_CacheEntry] = None
        if len(self._data) >= self._max:
            _, evicted = self._data.popitem(last=False)  # evict LRU (oldest)
            logger.debug(
                "WorkingMemory LRU evict: session=%s key=%s",
                entry.session_id,
                evicted.key,
            )

        self._data[entry.key] = entry
        return evicted

    def evict(self, key: str) -> Optional[_CacheEntry]:
        """Remove and return entry by key, or ``None`` if absent."""
        return self._data.pop(key, None)

    def all_entries(self) -> List[_CacheEntry]:
        """Return all entries in insertion order."""
        return list(self._data.values())

    def clear(self) -> int:
        """Remove all entries and return the count evicted."""
        n = len(self._data)
        self._data.clear()
        return n


# ---------------------------------------------------------------------------
# Public cache class
# ---------------------------------------------------------------------------

class WorkingMemoryCache:
    """Session-scoped, asyncio-safe working memory cache.

    One instance is typically shared across all sessions in a process.
    Session isolation is enforced by keying all internal structures on
    ``session_id``.

    Usage::

        wm = WorkingMemoryCache(max_items=100, default_ttl_seconds=1800)
        await wm.put("task", {"step": 3}, category="task_state", session_id="s1")
        entry = await wm.get("task", session_id="s1")
        # entry.value == {"step": 3}

    Args:
        max_items: Per-session item cap. When a session reaches this limit the
                   least-recently-used entry is evicted automatically.
        default_ttl_seconds: Default time-to-live in seconds. Pass ttl=None
                              to ``put()`` to override per-entry. Pass 0 to
                              disable TTL globally.
        enabled: When False all operations are no-ops (returns None / empty).
    """

    def __init__(
        self,
        max_items: int = 100,
        default_ttl_seconds: int = 1800,
        enabled: bool = True,
    ) -> None:
        # Audit fix: reject invalid max_items (0 would crash _SessionLRU.put)
        if max_items < 1:
            raise ValueError(
                f"max_items must be >= 1, got {max_items}"
            )
        # Audit fix: reject negative TTL (would silently create always-expired entries)
        if default_ttl_seconds < 0:
            raise ValueError(
                f"default_ttl_seconds must be >= 0, got {default_ttl_seconds}"
            )
        self._max_items = max_items
        self._default_ttl = default_ttl_seconds
        self._enabled = enabled
        self._lock = asyncio.Lock()
        # session_id → _SessionLRU
        self._sessions: Dict[str, _SessionLRU] = {}
        # stats counters
        self._hits = 0
        self._misses = 0
        self._evictions = 0
        self._expirations = 0

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

    @property
    def enabled(self) -> bool:
        """Whether working memory is active (``False`` = all ops are no-ops)."""
        return self._enabled

    async def put(
        self,
        key: str,
        value: Any,
        category: str,
        *,
        session_id: str,
        ttl: Optional[int] = None,
    ) -> None:
        """Store a value under ``key`` in ``session_id``'s working memory.

        If the session is at capacity, the least-recently-used entry is
        evicted first to make room.

        Args:
            key: Unique string key within this session.
            value: Arbitrary payload to store.
            category: Must be one of the eight WorkingMemoryCategory values
                      (or their string equivalents).
            session_id: Session scope for isolation.
            ttl: Override default TTL in seconds.
                 Pass ``0`` to disable expiry for this entry.
                 Pass ``None`` to use ``default_ttl_seconds``.
        """
        if not self._enabled:
            return

        _validate_category(category)

        # Audit fix: reject negative per-entry TTL
        if ttl is not None and ttl < 0:
            raise ValueError(f"ttl must be >= 0, got {ttl}")

        now = time.monotonic()
        effective_ttl = ttl if ttl is not None else self._default_ttl
        expires_at = (now + effective_ttl) if effective_ttl > 0 else None

        entry = _CacheEntry(
            key=key,
            value=value,
            category=category,
            session_id=session_id,
            created_at=now,
            expires_at=expires_at,
            last_accessed=now,
        )

        async with self._lock:
            lru = self._sessions.setdefault(session_id, _SessionLRU(self._max_items))
            evicted = lru.put(entry)
            if evicted is not None:
                self._evictions += 1

        logger.debug("WorkingMemory put: session=%s key=%s category=%s", session_id, key, category)

    async def get(self, key: str, *, session_id: str) -> Optional[_CacheEntry]:
        """Retrieve an entry by key. Returns None on miss or TTL expiry.

        Updates last_accessed on hit (LRU promotion).

        Args:
            key: The key passed to ``put()``.
            session_id: Session scope.

        Returns:
            The ``_CacheEntry`` (including ``.value``) or ``None``.
        """
        if not self._enabled:
            return None

        async with self._lock:
            lru = self._sessions.get(session_id)
            if lru is None:
                self._misses += 1
                return None

            entry = lru.get(key)
            if entry is None:
                self._misses += 1
                return None

            if _is_expired(entry):
                lru.evict(key)
                self._expirations += 1
                self._misses += 1
                logger.debug("WorkingMemory TTL expired: session=%s key=%s", session_id, key)
                return None

            entry.last_accessed = time.monotonic()
            self._hits += 1
            return entry

    async def get_value(self, key: str, *, session_id: str) -> Any:
        """Convenience wrapper: returns entry.value or None on miss."""
        entry = await self.get(key, session_id=session_id)
        return entry.value if entry is not None else None

    async def get_by_category(
        self,
        category: str,
        *,
        session_id: str,
    ) -> List[_CacheEntry]:
        """Return all non-expired entries matching ``category`` in ``session_id``.

        Args:
            category: One of the eight WorkingMemoryCategory values.
            session_id: Session scope.

        Returns:
            List of matching ``_CacheEntry`` objects, ordered by last_accessed
            descending (most recently used first).
        """
        if not self._enabled:
            return []

        _validate_category(category)

        async with self._lock:
            lru = self._sessions.get(session_id)
            if lru is None:
                return []

            now = time.monotonic()
            result: List[_CacheEntry] = []
            expired_keys: List[str] = []

            for entry in lru.all_entries():
                if entry.category != category:
                    continue
                if _is_expired(entry, now):
                    expired_keys.append(entry.key)
                    self._expirations += 1
                else:
                    result.append(entry)

            for k in expired_keys:
                lru.evict(k)

            result.sort(key=lambda e: e.last_accessed, reverse=True)
            return result

    async def evict(self, key: str, *, session_id: str) -> bool:
        """Explicitly remove an entry.

        Args:
            key: Key to evict.
            session_id: Session scope.

        Returns:
            True if the entry existed and was removed; False if not found.
        """
        if not self._enabled:
            return False

        async with self._lock:
            lru = self._sessions.get(session_id)
            if lru is None:
                return False
            removed = lru.evict(key)
            if removed:
                self._evictions += 1
            return removed is not None

    async def clear_session(self, session_id: str) -> int:
        """Remove ALL entries for ``session_id`` (e.g. on session end).

        Args:
            session_id: The session to clear.

        Returns:
            Number of entries removed.
        """
        if not self._enabled:
            return 0

        async with self._lock:
            lru = self._sessions.pop(session_id, None)
            if lru is None:
                return 0
            n = lru.clear()
            logger.debug("WorkingMemory clear_session: session=%s removed=%d", session_id, n)
            return n

    async def stats(self, session_id: Optional[str] = None) -> Dict[str, Any]:
        """Return cache statistics.

        Args:
            session_id: When provided, include per-session item counts and
                        category breakdown. When None, return global stats only.

        Returns:
            Dict with keys: enabled, hits, misses, hit_rate, evictions,
            expirations, total_sessions, and optionally session_items /
            session_categories.
        """
        async with self._lock:
            total_calls = self._hits + self._misses
            hit_rate = self._hits / total_calls if total_calls > 0 else 0.0

            result: Dict[str, Any] = {
                "enabled": self._enabled,
                "hits": self._hits,
                "misses": self._misses,
                "hit_rate": round(hit_rate, 4),
                "evictions": self._evictions,
                "expirations": self._expirations,
                "total_sessions": len(self._sessions),
                "max_items_per_session": self._max_items,
                "default_ttl_seconds": self._default_ttl,
            }

            if session_id is not None:
                lru = self._sessions.get(session_id)
                if lru is not None:
                    now = time.monotonic()
                    entries = lru.all_entries()
                    active = [e for e in entries if not _is_expired(e, now)]
                    by_cat: Dict[str, int] = {}
                    for e in active:
                        by_cat[e.category] = by_cat.get(e.category, 0) + 1
                    result["session_id"] = session_id
                    result["session_items"] = len(active)
                    result["session_categories"] = by_cat
                else:
                    result["session_id"] = session_id
                    result["session_items"] = 0
                    result["session_categories"] = {}

            return result

    # ------------------------------------------------------------------
    # Retrieval integration helpers
    # ------------------------------------------------------------------

    async def query(
        self,
        query_text: str,
        *,
        session_id: str,
        categories: Optional[List[str]] = None,
    ) -> List[_CacheEntry]:
        """Simple substring-match query over working memory for retrieval integration.

        Checks all (non-expired) entries in the session for the query text
        appearing in the string representation of the value. Returns matches
        ordered by last_accessed (most recent first).

        This is intentionally simple — working memory is small enough that
        a linear scan is fine (<100 items by default). No embeddings needed.

        Args:
            query_text: Text to search for (case-insensitive substring match).
            session_id: Session scope.
            categories: Optional allowlist of categories to search.
                        None = search all categories.

        Returns:
            Matching entries, ordered by recency.
        """
        if not self._enabled:
            return []

        query_lower = query_text.lower()

        async with self._lock:
            lru = self._sessions.get(session_id)
            if lru is None:
                return []

            now = time.monotonic()
            matches: List[_CacheEntry] = []
            expired_keys: List[str] = []

            for entry in lru.all_entries():
                if _is_expired(entry, now):
                    expired_keys.append(entry.key)
                    self._expirations += 1
                    continue
                if categories and entry.category not in categories:
                    continue
                # Substring match against string repr of value
                if query_lower in str(entry.value).lower() or query_lower in entry.key.lower():
                    matches.append(entry)

            for k in expired_keys:
                lru.evict(k)

        matches.sort(key=lambda e: e.last_accessed, reverse=True)
        return matches


# ---------------------------------------------------------------------------
# Private helpers
# ---------------------------------------------------------------------------

def _is_expired(entry: _CacheEntry, now: Optional[float] = None) -> bool:
    """Return True if the entry has passed its expiry time."""
    if entry.expires_at is None:
        return False
    t = now if now is not None else time.monotonic()
    return t >= entry.expires_at


_VALID_CATEGORIES = WorkingMemoryCategory.values()


def _validate_category(category: str) -> None:
    if category not in _VALID_CATEGORIES:
        raise ValueError(
            f"Invalid working memory category {category!r}. "
            f"Valid categories: {sorted(_VALID_CATEGORIES)}"
        )
