"""
tests/test_re_retrieval.py — Tests for mid-turn re-retrieval engine (#372).

All tests are self-contained with mocked retrieval pipeline.
No external dependencies, no DB, no LLM.
"""

from __future__ import annotations

import asyncio
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Set
from unittest.mock import AsyncMock, MagicMock, patch

import pytest

from cortex.core.re_retrieval import (
    ReRetrievalEngine,
    ReRetrievalRequest,
    ReRetrievalResult,
    ReRetrievalTrigger,
)


# ---------------------------------------------------------------------------
# Helpers / Fakes
# ---------------------------------------------------------------------------

@dataclass
class _FakeRetrievedItem:
    """Minimal stand-in for cortex.core.retrieval.RetrievedItem."""
    item_type: str = "claim"
    item_id: str = ""
    content: str = ""
    score: float = 0.9
    provenance: Optional[str] = None


@dataclass
class _FakeRetrievalResult:
    """Minimal stand-in for cortex.core.retrieval.RetrievalResult."""
    items: List[_FakeRetrievedItem] = field(default_factory=list)
    retrieval_id: str = "test-retrieval-id"
    mode: str = "lightweight"
    latency_ms: int = 50
    tokens_used: int = 200


class _FakeConfig:
    """Minimal config stub with re_retrieval defaults."""

    re_retrieval_enabled: bool = True
    re_retrieval_max_per_turn: int = 2
    re_retrieval_topic_shift_threshold: float = 0.7
    re_retrieval_min_budget_tokens: int = 500
    owner_id: str = "test-owner"
    retrieval_token_budget: int = 3000


class _FakeWorkingMemory:
    """Minimal async working memory stub."""

    def __init__(self) -> None:
        self._store: Dict[str, Any] = {}

    async def put(self, key: str, value: Any, category: str, *, session_id: str, ttl: Optional[int] = None) -> None:
        self._store[f"{session_id}:{key}"] = value

    async def get(self, key: str, *, session_id: str) -> Any:
        return self._store.get(f"{session_id}:{key}")


def _make_engine(
    retrieval_items: Optional[List[_FakeRetrievedItem]] = None,
    config: Optional[_FakeConfig] = None,
    working_memory: Optional[_FakeWorkingMemory] = None,
) -> ReRetrievalEngine:
    """Build a ReRetrievalEngine with mocked retrieval pipeline."""
    mock_retrieval = AsyncMock()
    mock_retrieval.retrieve = AsyncMock(
        return_value=_FakeRetrievalResult(items=retrieval_items or [])
    )

    cfg = config or _FakeConfig()
    wm = working_memory or _FakeWorkingMemory()

    return ReRetrievalEngine(
        retrieval=mock_retrieval,
        working_memory=wm,
        config=cfg,
    )


# ---------------------------------------------------------------------------
# Trigger detection tests
# ---------------------------------------------------------------------------

class TestShouldTrigger:
    """Tests for ReRetrievalEngine.should_trigger()."""

    @pytest.mark.asyncio
    async def test_new_entities_trigger_tool_result(self):
        """New entity names not in current context → TOOL_RESULT trigger."""
        engine = _make_engine()
        trigger = await engine.should_trigger(
            new_context="The results mention John Smith and Project Alpha",
            current_context=["We were discussing the weather"],
        )
        assert trigger == ReRetrievalTrigger.TOOL_RESULT

    @pytest.mark.asyncio
    async def test_topic_shift_trigger(self):
        """High keyword divergence → TOPIC_SHIFT trigger."""
        engine = _make_engine()
        # Force a low threshold so divergent text triggers
        engine._config.re_retrieval_topic_shift_threshold = 0.3
        trigger = await engine.should_trigger(
            new_context="quantum computing algorithms optimization neural networks",
            current_context=["cooking recipes pasta tomato sauce ingredients"],
        )
        assert trigger in (ReRetrievalTrigger.TOPIC_SHIFT, ReRetrievalTrigger.TOOL_RESULT)

    @pytest.mark.asyncio
    async def test_no_change_returns_none(self):
        """Same topic, no new entities → None."""
        engine = _make_engine()
        trigger = await engine.should_trigger(
            new_context="the weather is nice today",
            current_context=["the weather has been nice lately"],
        )
        assert trigger is None

    @pytest.mark.asyncio
    async def test_empty_new_context_returns_none(self):
        """Empty new_context → None."""
        engine = _make_engine()
        assert await engine.should_trigger("", ["some context"]) is None
        assert await engine.should_trigger("   ", ["some context"]) is None

    @pytest.mark.asyncio
    async def test_none_handled_gracefully(self):
        """None-like inputs don't crash."""
        engine = _make_engine()
        assert await engine.should_trigger("", []) is None

    @pytest.mark.asyncio
    async def test_empty_current_context(self):
        """Empty current_context with entities in new_context → trigger."""
        engine = _make_engine()
        trigger = await engine.should_trigger(
            new_context="Meeting with Alice Johnson about Cortex",
            current_context=[],
        )
        assert trigger == ReRetrievalTrigger.TOOL_RESULT


# ---------------------------------------------------------------------------
# Entity extraction tests
# ---------------------------------------------------------------------------

class TestEntityExtraction:
    """Tests for ReRetrievalEngine._extract_entities()."""

    def test_capitalized_words(self):
        engine = _make_engine()
        entities = engine._extract_entities("Alice went to Paris with Bob")
        assert "Alice" in entities
        assert "Paris" in entities
        assert "Bob" in entities

    def test_compound_names(self):
        engine = _make_engine()
        entities = engine._extract_entities("John Smith met Sarah Connor")
        assert "John Smith" in entities
        assert "Sarah Connor" in entities

    def test_acronyms(self):
        engine = _make_engine()
        entities = engine._extract_entities("The API uses NASA data and UNESCO reports")
        assert "API" in entities
        assert "NASA" in entities
        assert "UNESCO" in entities

    def test_stop_words_excluded(self):
        engine = _make_engine()
        entities = engine._extract_entities("The quick brown fox")
        # "The" is a stop word — should not appear
        assert "The" not in entities

    def test_empty_input(self):
        engine = _make_engine()
        assert engine._extract_entities("") == set()
        assert engine._extract_entities(None) == set()

    def test_short_names_excluded(self):
        """Names <= 2 chars should be excluded."""
        engine = _make_engine()
        entities = engine._extract_entities("Go to Al and see")
        assert "Al" not in entities


# ---------------------------------------------------------------------------
# Topic shift detection tests
# ---------------------------------------------------------------------------

class TestTopicShiftDetection:
    """Tests for ReRetrievalEngine._detect_topic_shift()."""

    def test_high_overlap_low_score(self):
        """Same topic → low divergence score."""
        engine = _make_engine()
        score = engine._detect_topic_shift(
            "python programming code functions",
            ["python code programming variables functions"],
        )
        assert score < 0.5  # high overlap → low divergence

    def test_low_overlap_high_score(self):
        """Different topics → high divergence score."""
        engine = _make_engine()
        score = engine._detect_topic_shift(
            "quantum physics particles electrons protons",
            ["cooking recipes pasta sauce ingredients tomato"],
        )
        assert score > 0.5  # low overlap → high divergence

    def test_empty_new_text(self):
        engine = _make_engine()
        assert engine._detect_topic_shift("", ["some context"]) == 0.0

    def test_empty_current_texts(self):
        engine = _make_engine()
        assert engine._detect_topic_shift("some text", []) == 0.0

    def test_identical_text(self):
        engine = _make_engine()
        score = engine._detect_topic_shift(
            "machine learning algorithms",
            ["machine learning algorithms"],
        )
        assert score == 0.0  # perfect overlap


# ---------------------------------------------------------------------------
# Retrieval tests
# ---------------------------------------------------------------------------

class TestRetrieve:
    """Tests for ReRetrievalEngine.retrieve()."""

    @pytest.mark.asyncio
    async def test_filters_existing_memories(self):
        """Already-injected memory IDs are filtered out."""
        items = [
            _FakeRetrievedItem(item_id="mem-1", content="memory one"),
            _FakeRetrievedItem(item_id="mem-2", content="memory two"),
            _FakeRetrievedItem(item_id="mem-3", content="memory three"),
        ]
        engine = _make_engine(retrieval_items=items)

        request = ReRetrievalRequest(
            trigger=ReRetrievalTrigger.TOOL_RESULT,
            new_context="some new information",
            session_id="s1",
            owner_id="test",
            existing_memory_ids=["mem-1", "mem-3"],
            budget_tokens=2000,
        )

        result = await engine.retrieve(request)
        assert len(result.items) == 1
        assert result.items[0].item_id == "mem-2"

    @pytest.mark.asyncio
    async def test_respects_token_budget(self):
        """Budget below minimum → empty result."""
        engine = _make_engine()
        engine._config.re_retrieval_min_budget_tokens = 500

        request = ReRetrievalRequest(
            trigger=ReRetrievalTrigger.TOOL_RESULT,
            new_context="some context",
            session_id="s1",
            owner_id="test",
            budget_tokens=100,  # below min
        )

        result = await engine.retrieve(request)
        assert len(result.items) == 0

    @pytest.mark.asyncio
    async def test_caches_results_in_working_memory(self):
        """Results are cached in working memory after retrieval."""
        items = [_FakeRetrievedItem(item_id="mem-1", content="cached memory")]
        wm = _FakeWorkingMemory()
        engine = _make_engine(retrieval_items=items, working_memory=wm)

        request = ReRetrievalRequest(
            trigger=ReRetrievalTrigger.TOOL_RESULT,
            new_context="trigger re-retrieval",
            session_id="s1",
            owner_id="test",
            budget_tokens=2000,
        )

        await engine.retrieve(request)
        cached = await wm.get("re_retrieval_0", session_id="s1")
        assert cached is not None
        assert cached["item_ids"] == ["mem-1"]

    @pytest.mark.asyncio
    async def test_max_per_turn_enforced(self):
        """After max_per_turn retrievals, subsequent calls return empty."""
        items = [_FakeRetrievedItem(item_id="mem-1", content="memory")]
        engine = _make_engine(retrieval_items=items)
        engine._config.re_retrieval_max_per_turn = 2

        request = ReRetrievalRequest(
            trigger=ReRetrievalTrigger.TOOL_RESULT,
            new_context="context",
            session_id="s1",
            owner_id="test",
            budget_tokens=2000,
        )

        # First two should succeed
        r1 = await engine.retrieve(request)
        assert len(r1.items) == 1
        r2 = await engine.retrieve(request)
        assert len(r2.items) == 1

        # Third should be blocked
        r3 = await engine.retrieve(request)
        assert len(r3.items) == 0

    @pytest.mark.asyncio
    async def test_min_budget_threshold(self):
        """Retrieval skipped when budget below min_budget_tokens."""
        engine = _make_engine()
        engine._config.re_retrieval_min_budget_tokens = 1000

        request = ReRetrievalRequest(
            trigger=ReRetrievalTrigger.TOOL_RESULT,
            new_context="context",
            session_id="s1",
            owner_id="test",
            budget_tokens=999,
        )

        result = await engine.retrieve(request)
        assert len(result.items) == 0

    @pytest.mark.asyncio
    async def test_latency_tracked(self):
        """Result includes latency_ms > 0."""
        items = [_FakeRetrievedItem(item_id="mem-1", content="memory")]
        engine = _make_engine(retrieval_items=items)

        request = ReRetrievalRequest(
            trigger=ReRetrievalTrigger.TOOL_RESULT,
            new_context="context",
            session_id="s1",
            owner_id="test",
            budget_tokens=2000,
        )

        result = await engine.retrieve(request)
        assert result.latency_ms >= 0

    @pytest.mark.asyncio
    async def test_empty_context_returns_empty(self):
        """Empty new_context in request → empty result."""
        engine = _make_engine()

        request = ReRetrievalRequest(
            trigger=ReRetrievalTrigger.TOOL_RESULT,
            new_context="",
            session_id="s1",
            owner_id="test",
            budget_tokens=2000,
        )

        result = await engine.retrieve(request)
        assert len(result.items) == 0


# ---------------------------------------------------------------------------
# Explicit request trigger tests
# ---------------------------------------------------------------------------

class TestExplicitRequest:
    """Tests for EXPLICIT_REQUEST trigger type."""

    @pytest.mark.asyncio
    async def test_explicit_request_always_fires(self):
        """EXPLICIT_REQUEST trigger bypasses detection — always retrieves."""
        items = [_FakeRetrievedItem(item_id="mem-1", content="memory")]
        engine = _make_engine(retrieval_items=items)

        # Explicit request directly calls retrieve(), not should_trigger()
        request = ReRetrievalRequest(
            trigger=ReRetrievalTrigger.EXPLICIT_REQUEST,
            new_context="re-retrieve for me",
            session_id="s1",
            owner_id="test",
            budget_tokens=2000,
        )

        result = await engine.retrieve(request)
        assert result.trigger == ReRetrievalTrigger.EXPLICIT_REQUEST
        assert len(result.items) == 1


# ---------------------------------------------------------------------------
# Turn counter reset tests
# ---------------------------------------------------------------------------

class TestTurnCounter:
    """Tests for turn counter management."""

    @pytest.mark.asyncio
    async def test_reset_counter_allows_new_retrievals(self):
        """reset_turn_counter() resets the per-turn limit."""
        items = [_FakeRetrievedItem(item_id="mem-1", content="memory")]
        engine = _make_engine(retrieval_items=items)
        engine._config.re_retrieval_max_per_turn = 1

        request = ReRetrievalRequest(
            trigger=ReRetrievalTrigger.TOOL_RESULT,
            new_context="context",
            session_id="s1",
            owner_id="test",
            budget_tokens=2000,
        )

        # First retrieval succeeds
        r1 = await engine.retrieve(request)
        assert len(r1.items) == 1

        # Second blocked
        r2 = await engine.retrieve(request)
        assert len(r2.items) == 0

        # Reset counter
        engine.reset_turn_counter()

        # Now succeeds again
        r3 = await engine.retrieve(request)
        assert len(r3.items) == 1


# ---------------------------------------------------------------------------
# Data type tests
# ---------------------------------------------------------------------------

class TestDataTypes:
    """Tests for ReRetrievalTrigger, ReRetrievalRequest, ReRetrievalResult."""

    def test_trigger_enum_values(self):
        assert ReRetrievalTrigger.TOOL_RESULT.value == "tool_result"
        assert ReRetrievalTrigger.TOPIC_SHIFT.value == "topic_shift"
        assert ReRetrievalTrigger.USER_CORRECTION.value == "user_correction"
        assert ReRetrievalTrigger.EXPLICIT_REQUEST.value == "explicit_request"

    def test_request_defaults(self):
        req = ReRetrievalRequest(
            trigger=ReRetrievalTrigger.TOOL_RESULT,
            new_context="text",
            session_id="s1",
            owner_id="o1",
        )
        assert req.existing_memory_ids == []
        assert req.budget_tokens == 2000

    def test_result_fields(self):
        result = ReRetrievalResult(
            items=[],
            trigger=ReRetrievalTrigger.TOOL_RESULT,
            query_used="test query",
            token_count=0,
            latency_ms=1.5,
        )
        assert result.query_used == "test query"
        assert result.latency_ms == 1.5
