"""
cortex/core/re_retrieval.py — Mid-Turn Re-Retrieval Engine (Issue #372).

When new information arrives mid-conversation (tool results, user corrections,
new entities mentioned), this module detects whether supplemental memory
retrieval should be triggered and executes it.

Trigger detection is fully deterministic (no LLM):
    - Entity extraction: capitalized words and compound names
    - Topic shift: keyword overlap ratio between new and current context

The actual retrieval delegates to RetrievalPipeline.retrieve() — no retrieval
logic is duplicated here.

Dependencies:
    - cortex.core.retrieval.RetrievalPipeline (M8) — actual retrieval
    - cortex.core.working_memory.WorkingMemoryCache (#371) — result caching
    - cortex.config.CortexConfig — re_retrieval.* settings
"""

from __future__ import annotations

import logging
import re
import time
from dataclasses import dataclass, field
from enum import Enum
from typing import TYPE_CHECKING, Dict, List, Optional, Set

if TYPE_CHECKING:
    from cortex.config import CortexConfig
    from cortex.core.retrieval import RetrievalPipeline, RetrievedItem
    from cortex.core.working_memory import WorkingMemoryCache

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Enums and data types
# ---------------------------------------------------------------------------

class ReRetrievalTrigger(str, Enum):
    """Reasons a mid-turn re-retrieval can fire."""

    TOOL_RESULT = "tool_result"
    TOPIC_SHIFT = "topic_shift"
    USER_CORRECTION = "user_correction"
    EXPLICIT_REQUEST = "explicit_request"


@dataclass
class ReRetrievalRequest:
    """Input for a supplemental retrieval.

    Attributes:
        trigger: Why this re-retrieval was requested.
        new_context: The new text/result that triggered re-retrieval.
        session_id: Active session.
        owner_id: Multi-tenancy scope key.
        existing_memory_ids: Already-injected memory IDs (to avoid duplicates).
        budget_tokens: Remaining token budget for supplemental results.
    """

    trigger: ReRetrievalTrigger
    new_context: str
    session_id: str
    owner_id: str
    existing_memory_ids: List[str] = field(default_factory=list)
    budget_tokens: int = 2000


@dataclass
class ReRetrievalResult:
    """Output from a supplemental retrieval.

    Attributes:
        items: New memories retrieved (already filtered for duplicates).
        trigger: Which trigger fired.
        query_used: The search query that was constructed.
        token_count: Total tokens in the returned items.
        latency_ms: End-to-end latency for the re-retrieval.
    """

    items: List  # List[RetrievedItem] — untyped to avoid circular import
    trigger: ReRetrievalTrigger
    query_used: str
    token_count: int
    latency_ms: float


# ---------------------------------------------------------------------------
# Stop words — excluded from entity extraction and topic analysis
# ---------------------------------------------------------------------------

_STOP_WORDS: frozenset[str] = frozenset({
    "the", "a", "an", "and", "or", "but", "in", "on", "at", "to", "for",
    "of", "with", "by", "from", "is", "it", "this", "that", "was", "are",
    "be", "been", "has", "have", "had", "do", "does", "did", "will", "would",
    "could", "should", "may", "might", "can", "shall", "not", "no", "if",
    "then", "else", "when", "where", "what", "which", "who", "whom", "how",
    "all", "each", "every", "both", "few", "more", "most", "other", "some",
    "such", "only", "own", "same", "so", "than", "too", "very", "just",
    "because", "as", "until", "while", "about", "between", "through",
    "during", "before", "after", "above", "below", "up", "down", "out",
    "off", "over", "under", "again", "further", "here", "there", "once",
    "also", "into", "its", "my", "your", "his", "her", "our", "their",
    "i", "me", "we", "you", "he", "she", "they", "them", "us",
    "none", "true", "false", "null", "error", "ok", "yes",
})

# Pattern for capitalized words (entity candidates)
_CAPITALIZED_WORD = re.compile(r"\b([A-Z][a-z]+(?:\s+[A-Z][a-z]+)*)\b")

# Pattern for ALLCAPS acronyms (2+ letters)
_ACRONYM = re.compile(r"\b([A-Z]{2,})\b")


# ---------------------------------------------------------------------------
# ReRetrievalEngine
# ---------------------------------------------------------------------------

class ReRetrievalEngine:
    """Detects when mid-turn re-retrieval is needed and executes it.

    Uses deterministic heuristics (entity extraction, keyword overlap) to
    decide whether new context warrants supplemental memory retrieval.
    Delegates actual retrieval to RetrievalPipeline.retrieve().

    Args:
        retrieval: The existing retrieval pipeline instance.
        working_memory: Working memory cache for result caching.
        config: CortexConfig with re_retrieval.* settings.
    """

    def __init__(
        self,
        retrieval: RetrievalPipeline,
        working_memory: WorkingMemoryCache,
        config: CortexConfig,
    ) -> None:
        self._retrieval = retrieval
        self._working_memory = working_memory
        self._config = config

        # Per-turn counter (reset by caller between turns)
        self._turn_retrieval_count: int = 0

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

    def reset_turn_counter(self) -> None:
        """Reset the per-turn re-retrieval counter. Call at the start of each turn."""
        self._turn_retrieval_count = 0

    async def should_trigger(
        self,
        new_context: str,
        current_context: List[str],
    ) -> Optional[ReRetrievalTrigger]:
        """Decide whether new_context warrants a supplemental retrieval.

        Detection order (first match wins):
            1. New entity names not present in current context → TOOL_RESULT
            2. Topic shift detected via keyword divergence → TOPIC_SHIFT

        Args:
            new_context: The new text (tool result, user message, etc.).
            current_context: List of text strings already in context.

        Returns:
            A ReRetrievalTrigger if re-retrieval should fire, or None.
        """
        if not new_context or not new_context.strip():
            return None

        # Check for new entities not in current context
        new_entities = self._extract_entities(new_context)
        if new_entities:
            current_text = " ".join(current_context) if current_context else ""
            current_entities = self._extract_entities(current_text)
            novel_entities = new_entities - current_entities
            if novel_entities:
                logger.debug(
                    "re_retrieval: new entities detected: %s",
                    novel_entities,
                )
                return ReRetrievalTrigger.TOOL_RESULT

        # Check for topic shift
        shift_score = self._detect_topic_shift(new_context, current_context)
        threshold = getattr(
            self._config, "re_retrieval_topic_shift_threshold", 0.7
        )
        if shift_score >= threshold:
            logger.debug(
                "re_retrieval: topic shift detected (score=%.3f, threshold=%.3f)",
                shift_score,
                threshold,
            )
            return ReRetrievalTrigger.TOPIC_SHIFT

        return None

    async def retrieve(self, request: ReRetrievalRequest) -> ReRetrievalResult:
        """Execute a supplemental retrieval based on a re-retrieval request.

        Steps:
            1. Check per-turn limit and minimum budget
            2. Build query from new_context
            3. Retrieve via RetrievalPipeline
            4. Filter out already-injected memories
            5. Cache results in working memory

        Args:
            request: ReRetrievalRequest with trigger, context, and constraints.

        Returns:
            ReRetrievalResult with new items and metadata.
        """
        t0 = time.monotonic()

        max_per_turn = getattr(self._config, "re_retrieval_max_per_turn", 2)
        min_budget = getattr(self._config, "re_retrieval_min_budget_tokens", 500)

        # Guard: max per-turn limit
        if self._turn_retrieval_count >= max_per_turn:
            logger.debug(
                "re_retrieval: skipping — max per-turn limit reached (%d/%d)",
                self._turn_retrieval_count,
                max_per_turn,
            )
            latency = (time.monotonic() - t0) * 1000
            return ReRetrievalResult(
                items=[],
                trigger=request.trigger,
                query_used="",
                token_count=0,
                latency_ms=latency,
            )

        # Guard: minimum budget
        if request.budget_tokens < min_budget:
            logger.debug(
                "re_retrieval: skipping — budget too low (%d < %d)",
                request.budget_tokens,
                min_budget,
            )
            latency = (time.monotonic() - t0) * 1000
            return ReRetrievalResult(
                items=[],
                trigger=request.trigger,
                query_used="",
                token_count=0,
                latency_ms=latency,
            )

        # Build query from new_context (use first 500 chars as search query)
        query = request.new_context[:500].strip()
        if not query:
            latency = (time.monotonic() - t0) * 1000
            return ReRetrievalResult(
                items=[],
                trigger=request.trigger,
                query_used="",
                token_count=0,
                latency_ms=latency,
            )

        # Import here to avoid circular imports at module level
        from cortex.core.retrieval import RetrievalConstraints

        # Retrieve via existing pipeline
        constraints = RetrievalConstraints(
            token_budget=request.budget_tokens,
            top_k=10,
        )

        retrieval_result = await self._retrieval.retrieve(
            query=query,
            mode="lightweight",
            constraints=constraints,
            owner_id=request.owner_id,
        )

        # Filter out already-injected memories
        existing_ids = set(request.existing_memory_ids)
        filtered_items = [
            item for item in retrieval_result.items
            if item.item_id not in existing_ids
        ]

        # Calculate token count from filtered items
        token_count = sum(
            len(item.content.split()) for item in filtered_items
        )

        # Cache results in working memory
        if filtered_items:
            try:
                cache_value = {
                    "trigger": request.trigger.value,
                    "query": query,
                    "item_ids": [item.item_id for item in filtered_items],
                }
                await self._working_memory.put(
                    key=f"re_retrieval_{self._turn_retrieval_count}",
                    value=cache_value,
                    category="tool_state",
                    session_id=request.session_id,
                )
            except Exception:
                logger.debug("re_retrieval: failed to cache results", exc_info=True)

        self._turn_retrieval_count += 1

        latency = (time.monotonic() - t0) * 1000
        return ReRetrievalResult(
            items=filtered_items,
            trigger=request.trigger,
            query_used=query,
            token_count=token_count,
            latency_ms=latency,
        )

    # ------------------------------------------------------------------
    # Deterministic helpers
    # ------------------------------------------------------------------

    def _extract_entities(self, text: str) -> Set[str]:
        """Extract entity names from text using deterministic rules.

        Entities are:
            - Capitalized words and compound names (e.g. "John Smith")
            - ALL-CAPS acronyms of 2+ letters (e.g. "API", "NASA")

        Excludes common stop words and single-character matches.

        Args:
            text: Input text to extract entities from.

        Returns:
            Set of entity name strings.
        """
        if not text:
            return set()

        entities: Set[str] = set()

        # Capitalized words and compound names
        for match in _CAPITALIZED_WORD.finditer(text):
            name = match.group(1)
            # Filter: skip single words that are stop words
            words = name.split()
            if len(words) == 1 and words[0].lower() in _STOP_WORDS:
                continue
            # Filter: skip very short names (1-2 chars)
            if len(name) <= 2:
                continue
            entities.add(name)

        # ALL-CAPS acronyms
        for match in _ACRONYM.finditer(text):
            acronym = match.group(1)
            if acronym.lower() not in _STOP_WORDS and len(acronym) >= 2:
                entities.add(acronym)

        return entities

    def _detect_topic_shift(
        self,
        new_text: str,
        current_texts: List[str],
    ) -> float:
        """Measure keyword divergence between new text and current context.

        Returns a float in [0.0, 1.0] where:
            0.0 = perfect overlap (same topic)
            1.0 = zero overlap (completely different topic)

        Uses lowercased, stop-word-filtered keyword sets and computes
        1 - Jaccard similarity.

        Args:
            new_text: The new text to compare.
            current_texts: List of existing context strings.

        Returns:
            Divergence score (0.0 = same topic, 1.0 = completely different).
        """
        if not new_text or not current_texts:
            return 0.0

        new_keywords = self._extract_keywords(new_text)
        current_keywords: Set[str] = set()
        for text in current_texts:
            current_keywords.update(self._extract_keywords(text))

        if not new_keywords or not current_keywords:
            return 0.0

        intersection = new_keywords & current_keywords
        union = new_keywords | current_keywords

        if not union:
            return 0.0

        jaccard = len(intersection) / len(union)
        return 1.0 - jaccard

    @staticmethod
    def _extract_keywords(text: str) -> Set[str]:
        """Extract meaningful keywords from text (lowercased, stop-words removed).

        Args:
            text: Input text.

        Returns:
            Set of keyword strings.
        """
        if not text:
            return set()

        words = re.findall(r"[a-zA-Z]+", text.lower())
        return {w for w in words if w not in _STOP_WORDS and len(w) > 2}
