"""In-memory TTL cache for anticipated context bundles.

Stores context bundles pushed proactively by the server via gRPC stream.
When the developer calls fetch(), the anticipation cache is checked first
to avoid a redundant round-trip when the server has already pre-fetched
the relevant context.

Lookup strategy:
  - Query present:  BM25 score between fetch query and cached prefetch
                    queries. Best match above threshold wins.
  - Query absent:   Scope (entity_id) match + freshest bundle wins.
"""

import logging
import time
from dataclasses import dataclass, field
from typing import Dict, List, Optional

from .bm25 import BM25, tokenize

logger = logging.getLogger("synap.sdk.cache.anticipation")

# Default BM25 score threshold — conservative to avoid false positives.
# Related short queries typically score 1.5-4.0; unrelated < 0.5.
_DEFAULT_BM25_THRESHOLD = 1.5


@dataclass
class _CacheEntry:
    """Internal cache entry with pre-tokenized metadata."""

    bundle: Dict
    search_queries: List[str]
    query_tokens: List[List[str]]  # pre-tokenized for BM25
    keyword_tokens: List[str]      # flattened keyword tokens (expanded vocabulary)
    entity_id: str                 # user_id / customer_id / "_client" / "_any"
    conversation_id: Optional[str] # conversation scope matching
    stored_at: float               # time.monotonic()
    bundle_type: str = "anticipation"  # "user_summary" or "anticipation"


class AnticipationCache:
    """In-memory TTL cache for anticipated context bundles.

    Dual-strategy lookup:
      1. Query present  → BM25-ranked matching against cached prefetch queries
      2. Query absent   → freshest scope-matched bundle
    """

    def __init__(
        self,
        ttl_seconds: int = 300,
        max_entries: int = 100,
        bm25_threshold: float = _DEFAULT_BM25_THRESHOLD,
    ):
        self._ttl = ttl_seconds
        self._max_entries = max_entries
        self._bm25_threshold = bm25_threshold
        self._entries: Dict[str, _CacheEntry] = {}  # bundle_id -> entry

    def store(self, bundle: Dict) -> None:
        """Store an anticipated bundle from the server.

        Args:
            bundle: ContextBundle dict (from proto conversion).
                    Must include ``search_queries`` list for BM25 matching.
        """
        self._evict_expired()

        # Evict oldest if at capacity
        if len(self._entries) >= self._max_entries:
            oldest_key = min(
                self._entries,
                key=lambda k: self._entries[k].stored_at,
            )
            del self._entries[oldest_key]

        search_queries = bundle.get("search_queries", [])
        query_tokens = [tokenize(q) for q in search_queries if q]

        # Tokenize expanded keywords (already individual words from server)
        search_keywords = bundle.get("search_keywords", [])
        keyword_tokens = tokenize(" ".join(search_keywords)) if search_keywords else []

        # Extract entity identifier from top-level anticipation metadata.
        # warm_cache injects _anticipation_user_id / _anticipation_customer_id.
        # ContextItem.entity_id is a graph reference, NOT the user/customer ID.
        entity_id = (
            bundle.get("_anticipation_user_id")
            or bundle.get("_anticipation_customer_id")
            or "_any"
        )
        conversation_id = bundle.get("_anticipation_conversation_id")
        bundle_type = bundle.get("_bundle_type", "anticipation")

        items_by_type = bundle.get("items_by_type", {})
        bundle_id = bundle.get("bundle_id", str(time.monotonic()))

        self._entries[bundle_id] = _CacheEntry(
            bundle=bundle,
            search_queries=search_queries,
            query_tokens=query_tokens,
            keyword_tokens=keyword_tokens,
            entity_id=entity_id,
            conversation_id=conversation_id,
            stored_at=time.monotonic(),
            bundle_type=bundle_type,
        )

        logger.info(
            "Anticipated bundle stored: bundle_id=%s type=%s queries=%s keywords=%d entity=%s conv=%s items=%d total_entries=%d",
            bundle_id,
            bundle_type,
            search_queries,
            len(keyword_tokens),
            entity_id,
            conversation_id,
            sum(len(v) for v in items_by_type.values() if isinstance(v, list)),
            len(self._entries),
        )

    def lookup(
        self,
        search_query: Optional[List[str]] = None,
        entity_id: Optional[str] = None,
        conversation_id: Optional[str] = None,
    ) -> Optional[Dict]:
        """Find a cached anticipation bundle matching the request.

        Args:
            search_query: Developer's fetch query (list of query strings).
                          If provided, uses BM25 scoring.
                          If None/empty, returns freshest scope-matched bundle.
            entity_id: Scope entity (user_id, customer_id, "_client").
                       Used as a filter in both strategies.
            conversation_id: Optional conversation filter. If provided,
                             entries with a matching conversation_id are preferred.

        Returns:
            Cached ContextBundle dict, or None if no match.
        """
        self._evict_expired()

        if not self._entries:
            logger.info(
                "Anticipation cache lookup: EMPTY cache (no entries). query=%s entity=%s conv=%s",
                search_query, entity_id, conversation_id,
            )
            return None

        # Filter by entity_id and/or conversation_id
        candidates = self._get_candidates(entity_id, conversation_id)
        if not candidates:
            logger.info(
                "Anticipation cache lookup: NO candidates after filter. "
                "total_entries=%d entity_filter=%s conv_filter=%s "
                "stored_entities=%s stored_convs=%s",
                len(self._entries),
                entity_id,
                conversation_id,
                [e.entity_id for e in self._entries.values()],
                [e.conversation_id for e in self._entries.values()],
            )
            return None

        # Choose strategy based on whether query is provided
        has_query = search_query and any(q.strip() for q in search_query if q)

        if has_query:
            return self._bm25_lookup(search_query, candidates)
        else:
            return self._freshness_lookup(candidates)

    def _get_candidates(
        self,
        entity_id: Optional[str],
        conversation_id: Optional[str] = None,
    ) -> Dict[str, _CacheEntry]:
        """Filter entries by entity_id and/or conversation_id."""
        results = {}
        for bid, entry in self._entries.items():
            # Entity filter
            if entity_id is not None and entry.entity_id not in (entity_id, "_any"):
                continue
            # Conversation filter — match if conversation_id matches OR entry has no conversation
            if conversation_id is not None:
                if entry.conversation_id and entry.conversation_id != conversation_id:
                    continue
            results[bid] = entry
        return results

    def _bm25_lookup(
        self,
        search_query: List[str],
        candidates: Dict[str, _CacheEntry],
    ) -> Optional[Dict]:
        """Case 1: Query present — BM25 ranked matching.

        The BM25 corpus for each entry includes:
        1. Tokenized prefetch queries (semantic query strings)
        2. Keyword tokens as an additional document (expanded user vocabulary)

        Keywords bridge the vocabulary gap between short LLM-generated
        prefetch queries and the developer's actual fetch query (which is
        often the user's full message). For example:
          prefetch_query = "user account details"
          keywords = ["account", "profile", "settings", "email", "password",
                      "login", "username", "subscription", "plan", "billing"]
          user message = "can you check my billing settings?"
          → "billing" and "settings" match keywords → BM25 score is high
        """
        # Tokenize the incoming query
        query_text = " ".join(q for q in search_query if q)
        query_tokens = tokenize(query_text)
        if not query_tokens:
            return self._freshness_lookup(candidates)

        best_bundle_id = None
        best_score = 0.0

        for bid, entry in candidates.items():
            if not entry.query_tokens and not entry.keyword_tokens:
                continue

            # Build corpus: prefetch queries + keywords as an extra document
            corpus = list(entry.query_tokens)
            if entry.keyword_tokens:
                corpus.append(entry.keyword_tokens)

            if not corpus:
                continue

            bm25 = BM25(corpus)
            entry_scores = bm25.scores(query_tokens)
            max_score = max(entry_scores) if entry_scores else 0.0

            if max_score > best_score:
                best_score = max_score
                best_bundle_id = bid

        # Adaptive threshold: short queries (1-4 tokens) produce lower BM25
        # scores since fewer terms can match. Scale threshold with query length
        # but cap at the configured maximum for long queries.
        effective_threshold = max(0.6, min(self._bm25_threshold, 0.3 * len(query_tokens)))

        if best_bundle_id and best_score >= effective_threshold:
            bundle = self._entries[best_bundle_id].bundle
            logger.info(
                "Anticipation cache HIT (BM25 score=%.2f, threshold=%.2f, query_tokens=%d, query=%s)",
                best_score,
                effective_threshold,
                len(query_tokens),
                search_query[:50] if search_query else None,
            )
            return bundle

        # Log details about the miss for diagnostics
        logger.info(
            "Anticipation cache MISS (best BM25 score=%.2f, threshold=%.2f, "
            "query_tokens=%d, query=%s, candidates=%d, candidate_queries=%s, candidate_keywords=%s)",
            best_score,
            effective_threshold,
            len(query_tokens),
            search_query[:80] if search_query else None,
            len(candidates),
            {bid: e.search_queries for bid, e in candidates.items()},
            {bid: e.keyword_tokens[:10] for bid, e in candidates.items()},
        )
        return None

    def _freshness_lookup(
        self,
        candidates: Dict[str, _CacheEntry],
    ) -> Optional[Dict]:
        """Case 2: No query — return freshest user_summary bundle only.

        Regular anticipation bundles should only match via BM25.
        User summary bundles are the only ones that make sense without
        a specific query (they represent broad user context).
        """
        if not candidates:
            return None

        summary_candidates = {
            bid: e for bid, e in candidates.items()
            if e.bundle_type == "user_summary"
        }
        if not summary_candidates:
            logger.debug("Anticipation cache: no user_summary bundles for freshness lookup")
            return None

        freshest_bid = max(summary_candidates, key=lambda k: summary_candidates[k].stored_at)
        bundle = summary_candidates[freshest_bid].bundle
        logger.debug("Anticipation cache HIT (freshest user_summary, no query)")
        return bundle

    def lookup_user_summary(
        self,
        entity_id: Optional[str] = None,
    ) -> Optional[Dict]:
        """Find the freshest user_summary bundle for the given entity.

        Used by the SDK's periodic injection logic to retrieve the user
        summary for mixing into fetch responses every N turns.
        """
        self._evict_expired()
        candidates = {
            bid: entry for bid, entry in self._entries.items()
            if entry.bundle_type == "user_summary"
            and (entity_id is None or entry.entity_id in (entity_id, "_any"))
        }
        if not candidates:
            return None
        freshest_bid = max(candidates, key=lambda k: candidates[k].stored_at)
        return candidates[freshest_bid].bundle

    def _evict_expired(self) -> None:
        """Remove expired entries."""
        now = time.monotonic()
        expired = [
            bid
            for bid, entry in self._entries.items()
            if now - entry.stored_at > self._ttl
        ]
        for bid in expired:
            del self._entries[bid]

    def clear(self) -> None:
        """Clear all cached entries."""
        self._entries.clear()
