"""
tests/test_retrieval_wm_session_and_filter.py — Regression tests for P8-1 and P8-2.

P8-1: working_memory_retrieve uses the real session_id, not owner_id.
P8-2: lightweight retrieval filters applied before top_k truncation.
"""

from __future__ import annotations

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

import pytest

from cortex.core.retrieval import RetrievalConstraints, RetrievedItem
from cortex.core.retrieval_search import working_memory_retrieve


# ---------------------------------------------------------------------------
# P8-1: WM retrieval with real session_id
# ---------------------------------------------------------------------------


class _FakeWMEntry:
    """Mimics a working-memory cache entry returned by WorkingMemoryCache.query()."""

    def __init__(self, key: str, value: Any, session_id: str, category: str = "task_state"):
        self.key = key
        self.value = value
        self.session_id = session_id
        self.category = category


class _FakeWorkingMemory:
    """Minimal stand-in for WorkingMemoryCache with session-scoped storage."""

    def __init__(self, enabled: bool = True):
        self.enabled = enabled
        self._entries: Dict[str, List[_FakeWMEntry]] = {}

    def add(self, entry: _FakeWMEntry):
        self._entries.setdefault(entry.session_id, []).append(entry)

    async def query(self, query: str, session_id: str) -> List[_FakeWMEntry]:
        return self._entries.get(session_id, [])


class TestWMSessionScoping:
    """P8-1: WM retrieval must use the real session_id, not owner_id."""

    @pytest.mark.asyncio
    async def test_wm_retrieve_uses_real_session_id(self):
        """WM entries stored under 'sess-42' are found when session_id='sess-42'."""
        wm = _FakeWorkingMemory()
        wm.add(_FakeWMEntry(key="goal", value="finish report", session_id="sess-42"))

        constraints = RetrievalConstraints()

        # With correct session_id → found
        items = await working_memory_retrieve(
            wm, "what is the goal?", constraints, owner_id="default", session_id="sess-42",
        )
        assert len(items) == 1
        assert "finish report" in items[0].content

    @pytest.mark.asyncio
    async def test_wm_retrieve_owner_id_misses_real_session(self):
        """Using owner_id as session_id misses entries stored under a real session."""
        wm = _FakeWorkingMemory()
        wm.add(_FakeWMEntry(key="goal", value="finish report", session_id="sess-42"))

        constraints = RetrievalConstraints()

        # With owner_id (no session_id) → falls back to owner_id, misses entries
        items = await working_memory_retrieve(
            wm, "what is the goal?", constraints, owner_id="default", session_id=None,
        )
        assert len(items) == 0

    @pytest.mark.asyncio
    async def test_wm_fallback_logs_warning(self, caplog):
        """When session_id is None, a WARNING is logged about falling back to owner_id."""
        wm = _FakeWorkingMemory()
        constraints = RetrievalConstraints()

        with caplog.at_level(logging.WARNING):
            await working_memory_retrieve(
                wm, "anything", constraints, owner_id="default", session_id=None,
            )

        assert any("falling back to owner_id" in r.message for r in caplog.records)

    @pytest.mark.asyncio
    async def test_wm_fallback_works_when_entries_stored_under_owner(self):
        """If entries are stored under the owner_id key and no session_id given,
        fallback still works (backward compat)."""
        wm = _FakeWorkingMemory()
        wm.add(_FakeWMEntry(key="x", value="val", session_id="default"))

        constraints = RetrievalConstraints()

        items = await working_memory_retrieve(
            wm, "query", constraints, owner_id="default", session_id=None,
        )
        assert len(items) == 1


# ---------------------------------------------------------------------------
# P8-2: Filters applied before top_k truncation
# ---------------------------------------------------------------------------


@dataclass
class _MockConfig:
    embedding_dim: int = 4
    retrieval_token_budget: int = 10000
    total_memory_token_budget: int = 40000
    retrieval_recency_weight: float = 0.0  # disable for simpler testing
    retrieval_temporal_boost_weight: float = 0.0


def _make_config(**overrides) -> _MockConfig:
    return _MockConfig(**overrides)


@dataclass
class _MockFTSResult:
    item_id: str
    content: str
    item_type: str = "claim"
    rank: float = -1.0


@dataclass
class _MockVectorResult:
    item_id: str
    score: float
    content: str
    item_type: str = "claim"


@dataclass
class _MockClaim:
    memory_class: str = "semantic"
    salience: str = "medium"
    status: str = "active"
    explicitness: str = "inferred"
    stability: str = "volatile"
    category: str = "misc"
    is_deleted: int = 0
    temporal: Optional[str] = None
    temporal_marker: Optional[str] = None
    event_date: Optional[str] = None
    cause: Optional[str] = None
    effect: Optional[str] = None
    state_before: Optional[str] = None
    state_after: Optional[str] = None
    entity_id: Optional[str] = None
    entity_type: str = ""
    scope: str = ""
    memory_type: Optional[str] = None
    mentions: int = 0
    last_mentioned_at: Optional[str] = None
    last_accessed: Optional[str] = None
    superseded_by: Optional[str] = None
    owner_id: str = "default"
    created_at: Optional[str] = None
    metadata_json: Optional[str] = None
    source_session: Optional[str] = None
    valid_to: Optional[str] = None
    sensitivity: Optional[str] = None


def _build_mock_storage(fts_results, vec_results, claim_cache):
    """Build a storage mock with all required async methods.

    NOTE: retrieval_lightweight.py calls get_claims_batch via
    ``getattr(type(storage), 'get_claims_batch', None)`` then invokes
    the unbound method directly (first positional arg = ids list).
    We therefore deliberately omit get_claims_batch as an instance method
    so that the code falls back to per-ID get_claim.
    """

    class MockStorage:
        async def fts_search(self, *a, **kw):
            return fts_results

        async def vector_search(self, *a, **kw):
            return vec_results

        async def get_claim(self, cid):
            return claim_cache.get(cid)

        async def get_claims_sensitivity_batch(self, ids):
            # Return all as "normal" sensitivity so nothing is excluded
            return {iid: "normal" for iid in ids}

    return MockStorage()


class TestFilterBeforeTruncation:
    """P8-2: Filters must not be starved by top_k truncation."""

    @pytest.mark.asyncio
    async def test_heavy_type_filter_still_returns_top_k(self):
        """With memory_types filtering out 60%+ of results, we should still
        get top_k items (assuming enough valid candidates exist)."""
        from cortex.core.retrieval_lightweight import lightweight_retrieve

        top_k = 5
        # Create 30 items: 12 'claim', 18 'event'
        # With memory_types=['claim'], only 12 are eligible.
        # Old code: fused[:5] → might get 0-2 claims out of 5.
        # New code: scans wider, should get 5 claims.
        all_items = []
        for i in range(30):
            item_type = "claim" if i % 5 < 2 else "event"  # ~40% claim, 60% event
            all_items.append((f"item-{i:03d}", item_type, f"content for item {i}"))

        fts_results = [
            _MockFTSResult(item_id=iid, content=content, item_type=itype, rank=-float(30 - idx))
            for idx, (iid, itype, content) in enumerate(all_items)
        ]
        vec_results = [
            _MockVectorResult(item_id=iid, score=1.0 - idx * 0.01, content=content, item_type=itype)
            for idx, (iid, itype, content) in enumerate(all_items)
        ]

        claim_cache = {
            iid: _MockClaim(owner_id="default")
            for iid, itype, _ in all_items
            if itype == "claim"
        }

        storage = _build_mock_storage(fts_results, vec_results, claim_cache)

        config = _make_config()

        constraints = RetrievalConstraints(
            top_k=top_k,
            memory_types=["claim"],  # filter out all events
        )

        async def mock_embed(text):
            return [0.1, 0.2, 0.3, 0.4]

        items, _info = await lightweight_retrieve(
            storage, config, "test query", constraints,
            owner_id="default", rrf_k=60, candidate_multiplier=3,
            embed_fn=mock_embed,
        )

        # Should have exactly top_k claims (there are 12 eligible, we want 5)
        assert len(items) == top_k
        assert all(it.item_type == "claim" for it in items)

    @pytest.mark.asyncio
    async def test_fewer_than_top_k_when_not_enough_candidates(self):
        """When there aren't enough valid candidates even in the wider window,
        we get as many as available (not crash)."""
        from cortex.core.retrieval_lightweight import lightweight_retrieve

        top_k = 10
        # Only 3 claims total, rest are events
        all_items = []
        for i in range(20):
            item_type = "claim" if i < 3 else "event"
            all_items.append((f"item-{i:03d}", item_type, f"content for item {i}"))

        fts_results = [
            _MockFTSResult(item_id=iid, content=content, item_type=itype, rank=-float(20 - idx))
            for idx, (iid, itype, content) in enumerate(all_items)
        ]
        vec_results = [
            _MockVectorResult(item_id=iid, score=1.0 - idx * 0.01, content=content, item_type=itype)
            for idx, (iid, itype, content) in enumerate(all_items)
        ]

        claim_cache = {
            iid: _MockClaim(owner_id="default")
            for iid, itype, _ in all_items
            if itype == "claim"
        }

        storage = _build_mock_storage(fts_results, vec_results, claim_cache)

        config = _make_config()

        constraints = RetrievalConstraints(
            top_k=top_k,
            memory_types=["claim"],
        )

        async def mock_embed(text):
            return [0.1, 0.2, 0.3, 0.4]

        items, _info = await lightweight_retrieve(
            storage, config, "test query", constraints,
            owner_id="default", rrf_k=60, candidate_multiplier=3,
            embed_fn=mock_embed,
        )

        # Only 3 claims exist, so we get 3 (not 10, not crash)
        assert len(items) == 3
        assert all(it.item_type == "claim" for it in items)
