"""
Tests for cortex/core/retrieval.py — M8 Basic Hybrid Retrieval Pipeline.

Coverage:
  - rrf_fusion math (correct RRF scores, k=60, rank ordering)
  - cosine_similarity (correct values, zero-vector guard, dim mismatch)
  - _normalise_bm25 (polarity inversion, [0,1] range)
  - RetrievalPipeline.retrieve() end-to-end with mock storage
  - BM25-only path (no embedding available)
  - Vector-only path (FTS5 raises)
  - Token budget enforcement (items trimmed when over budget)
  - retrieval_log called on every retrieve()
  - assemble_context() respects budget
  - Status penalty (#574) — superseded/deprecated claims score lower
"""

from __future__ import annotations

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

import pytest

from cortex.core.retrieval import (
    RetrievalConstraints,
    RetrievalPipeline,
    RetrievalResult,
    RetrievedItem,
    _FTSResult,
    _VectorResult,
    _normalise_bm25,
    cosine_similarity,
    rrf_fusion,
)


# ---------------------------------------------------------------------------
# Fixtures & helpers
# ---------------------------------------------------------------------------

@dataclass
class MockConfig:
    retrieval_token_budget: int = 500
    total_memory_token_budget: int = 4000
    cornerstone_token_budget: int = 800
    embedding_dim: int = 4
    embedding_model: str = "mock"


@dataclass
class MockFTSResult:
    item_id: str
    content: str
    item_type: str = "claim"
    rank: float = -1.0  # FTS5 bm25() is negative


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


def make_storage(
    fts_items: Optional[List[MockFTSResult]] = None,
    vector_items: Optional[List[MockVectorResult]] = None,
    fts_raises: Optional[Exception] = None,
    vector_raises: Optional[Exception] = None,
) -> MagicMock:
    storage = MagicMock()
    storage.embed = AsyncMock(return_value=[0.1, 0.2, 0.3, 0.4])

    # Route FTS items by table to preserve item_type in content_map
    _claims = [i for i in (fts_items or []) if getattr(i, "item_type", "claim") == "claim"]
    _events = [i for i in (fts_items or []) if getattr(i, "item_type", "claim") == "event"]

    if fts_raises:
        storage.fts_search = AsyncMock(side_effect=fts_raises)
    else:
        async def _fts(query, table="claims", top_k=50, owner_id="default"):
            return _claims if table == "claims" else _events
        storage.fts_search = AsyncMock(side_effect=_fts)

    if vector_raises:
        storage.vector_search = AsyncMock(side_effect=vector_raises)
    else:
        storage.vector_search = AsyncMock(return_value=vector_items or [])

    storage.log_retrieval = AsyncMock()

    # Default: all claim items get sensitivity="normal"
    _sens = {i.item_id: "normal" for i in (fts_items or []) + (vector_items or [])
             if getattr(i, "item_type", "claim") == "claim"}
    storage.get_claims_sensitivity_batch = AsyncMock(return_value=_sens)

    return storage


def make_token_budget(tokens_per_call: int = 10) -> MagicMock:
    tb = MagicMock()
    tb.count_tokens = MagicMock(return_value=tokens_per_call)
    budget_obj = MagicMock()
    budget_obj.remaining = 200
    tb.create_budget = MagicMock(return_value=budget_obj)
    return tb


def make_pipeline(**kwargs) -> RetrievalPipeline:
    storage = kwargs.pop("storage", make_storage())
    config = kwargs.pop("config", MockConfig())
    token_budget = kwargs.pop("token_budget", make_token_budget())
    return RetrievalPipeline(storage, config, token_budget, **kwargs)


# ---------------------------------------------------------------------------
# Unit tests: rrf_fusion
# ---------------------------------------------------------------------------

class TestRRFFusion:
    def test_single_list_scores(self):
        """RRF score for rank 0 with k=60 should be 1/(60+0+1) = 1/61."""
        items = [("a", 1.0), ("b", 0.5)]
        result = rrf_fusion([items], k=60)
        scores = dict(result)
        assert abs(scores["a"] - 1.0 / 61) < 1e-9
        assert abs(scores["b"] - 1.0 / 62) < 1e-9

    def test_two_lists_additive(self):
        """Item appearing in both lists gets higher score than item in one list."""
        list1 = [("a", 1.0), ("b", 0.5)]
        list2 = [("a", 0.9), ("c", 0.3)]
        result = rrf_fusion([list1, list2], k=60)
        scores = dict(result)
        # "a" appears in both → higher than "b" and "c" which appear in one
        assert scores["a"] > scores["b"]
        assert scores["a"] > scores["c"]

    def test_order_descending(self):
        """Result is sorted best-first."""
        items = [("a", 1.0), ("b", 0.9), ("c", 0.8)]
        result = rrf_fusion([items], k=60)
        scores = [s for _, s in result]
        assert scores == sorted(scores, reverse=True)

    def test_empty_lists(self):
        assert rrf_fusion([]) == []
        assert rrf_fusion([[]]) == []

    def test_k_parameter(self):
        """Larger k → smaller score differences across ranks."""
        items = [("a", 1.0), ("b", 0.0)]
        r60 = dict(rrf_fusion([items], k=60))
        r10 = dict(rrf_fusion([items], k=10))
        diff_60 = r60["a"] - r60["b"]
        diff_10 = r10["a"] - r10["b"]
        assert diff_10 > diff_60, "Smaller k should amplify rank differences"

    def test_rrf_formula_explicit(self):
        """Manually verify two-list, two-item RRF with k=60."""
        list1 = [("x", 10.0), ("y", 5.0)]  # x@rank0, y@rank1
        list2 = [("y", 8.0), ("x", 3.0)]  # y@rank0, x@rank1
        result = dict(rrf_fusion([list1, list2], k=60))
        # x: 1/61 + 1/62
        # y: 1/62 + 1/61  — same!
        assert abs(result["x"] - result["y"]) < 1e-9

    def test_deduplication(self):
        """Each item_id should appear exactly once in output."""
        list1 = [("a", 1.0)] * 3  # duplicates in input (malformed but guard it)
        # rrf_fusion accumulates, so "a" appears once in output with sum
        result = rrf_fusion([list1], k=60)
        ids = [item_id for item_id, _ in result]
        # With duplicates: each rank contributes separately — but output should have unique ids
        # Our implementation doesn't dedup input ranks, but scores dict key is unique
        assert len(ids) == len(set(ids))


# ---------------------------------------------------------------------------
# Unit tests: cosine_similarity
# ---------------------------------------------------------------------------

class TestCosineSimilarity:
    def test_identical_vectors(self):
        v = [1.0, 0.0, 0.0]
        assert abs(cosine_similarity(v, v) - 1.0) < 1e-9

    def test_orthogonal_vectors(self):
        a = [1.0, 0.0]
        b = [0.0, 1.0]
        assert abs(cosine_similarity(a, b)) < 1e-9

    def test_opposite_vectors(self):
        a = [1.0, 0.0]
        b = [-1.0, 0.0]
        assert abs(cosine_similarity(a, b) - (-1.0)) < 1e-9

    def test_zero_vector(self):
        a = [0.0, 0.0]
        b = [1.0, 1.0]
        assert cosine_similarity(a, b) == 0.0

    def test_dimension_mismatch(self):
        with pytest.raises(ValueError, match="dimension mismatch"):
            cosine_similarity([1.0, 2.0], [1.0])

    def test_known_value(self):
        a = [3.0, 4.0]
        b = [4.0, 3.0]
        # dot=24, |a|=5, |b|=5 → 24/25=0.96
        assert abs(cosine_similarity(a, b) - 0.96) < 1e-9


# ---------------------------------------------------------------------------
# Unit tests: _normalise_bm25
# ---------------------------------------------------------------------------

class TestNormaliseBM25:
    def test_empty(self):
        assert _normalise_bm25([]) == []

    def test_single_item(self):
        results = [_FTSResult(item_type="claim", item_id="x", content="foo", rank=-2.0)]
        norm = _normalise_bm25(results)
        assert len(norm) == 1
        # single item — after negation rank=2, min==max → span=1 → (2-2)/1=0
        assert norm[0][0] == "x"
        assert norm[0][1] == 0.0

    def test_two_items_range(self):
        results = [
            _FTSResult(item_type="claim", item_id="best", content="", rank=-5.0),   # best BM25
            _FTSResult(item_type="claim", item_id="worst", content="", rank=-1.0),  # worst BM25
        ]
        norm = dict(_normalise_bm25(results))
        # After negation: best=5, worst=1. min=1, max=5, span=4
        # best: (5-1)/4 = 1.0; worst: (1-1)/4 = 0.0
        assert abs(norm["best"] - 1.0) < 1e-9
        assert abs(norm["worst"] - 0.0) < 1e-9


# ---------------------------------------------------------------------------
# Integration tests: RetrievalPipeline
# ---------------------------------------------------------------------------

class TestRetrievalPipeline:
    @pytest.mark.asyncio
    async def test_basic_hybrid_returns_results(self):
        """Full pipeline with both BM25 and vector results."""
        fts = [MockFTSResult("c1", "memory about cats", rank=-3.0),
               MockFTSResult("c2", "memory about dogs", rank=-1.5)]
        vec = [MockVectorResult("c1", "memory about cats", score=0.9),
               MockVectorResult("c3", "memory about fish", score=0.7)]
        storage = make_storage(fts_items=fts, vector_items=vec)
        pipeline = make_pipeline(storage=storage)

        result = await pipeline.retrieve("cats", owner_id="test")

        assert isinstance(result, RetrievalResult)
        assert len(result.items) > 0
        assert result.mode == "lightweight"
        assert result.latency_ms >= 0
        # c1 should score high (in both lists)
        ids = [it.item_id for it in result.items]
        assert "c1" in ids

    @pytest.mark.asyncio
    async def test_retrieval_log_always_called(self):
        """log_retrieval must be called on every retrieve()."""
        storage = make_storage(
            fts_items=[MockFTSResult("a", "test content", rank=-2.0)],
            vector_items=[MockVectorResult("a", "test content", score=0.8)],
        )
        pipeline = make_pipeline(storage=storage)
        await pipeline.retrieve("test query")

        storage.log_retrieval.assert_called_once()
        call_kwargs = storage.log_retrieval.call_args.kwargs
        assert "query_text" in call_kwargs
        assert call_kwargs["query_text"] == "test query"
        assert "mode" in call_kwargs
        assert "latency_ms" in call_kwargs

    @pytest.mark.asyncio
    async def test_log_called_even_on_empty_results(self):
        """Even with no results, retrieval_log must be called."""
        storage = make_storage(fts_items=[], vector_items=[])
        pipeline = make_pipeline(storage=storage)
        await pipeline.retrieve("nothing")
        storage.log_retrieval.assert_called_once()

    @pytest.mark.asyncio
    async def test_bm25_only_when_no_embedding(self):
        """When embedding is unavailable, BM25 results still return."""
        fts = [MockFTSResult("b1", "keyword content", rank=-4.0),
               MockFTSResult("b2", "other content", rank=-1.0)]
        storage = make_storage(fts_items=fts, vector_items=[])
        storage.embed = AsyncMock(return_value=None)  # no embedding
        storage.vector_search = AsyncMock(return_value=[])
        pipeline = make_pipeline(storage=storage)
        # Override _embed to return None
        pipeline._embed = AsyncMock(return_value=None)  # type: ignore

        result = await pipeline.retrieve("keyword", owner_id="u1")
        assert len(result.items) >= 1  # BM25 results present
        assert all(it.item_id in ("b1", "b2") for it in result.items)

    @pytest.mark.asyncio
    async def test_vector_only_when_fts_fails(self):
        """When FTS raises, vector results still return."""
        vec = [MockVectorResult("v1", "semantic content", score=0.95)]
        storage = make_storage(
            vector_items=vec,
            fts_raises=Exception("FTS5 unavailable"),
        )
        pipeline = make_pipeline(storage=storage)

        result = await pipeline.retrieve("semantic query")
        # Should still have vector result
        assert any(it.item_id == "v1" for it in result.items)

    @pytest.mark.asyncio
    async def test_token_budget_enforced(self):
        """Items exceeding token budget are trimmed."""
        # Each item costs 100 tokens; budget is 250 → max 2 items
        fts = [MockFTSResult(f"c{i}", f"content item {i}", rank=float(-i))
               for i in range(1, 6)]
        vec = [MockVectorResult(f"c{i}", f"content item {i}", score=1.0 / i)
               for i in range(1, 6)]
        storage = make_storage(fts_items=fts, vector_items=vec)

        tb = make_token_budget(tokens_per_call=100)
        config = MockConfig(retrieval_token_budget=250)
        pipeline = make_pipeline(storage=storage, config=config, token_budget=tb)

        result = await pipeline.retrieve("query", constraints=RetrievalConstraints(top_k=10))
        assert len(result.items) <= 2
        assert result.tokens_used <= 250

    @pytest.mark.asyncio
    async def test_token_budget_from_constraints_overrides_config(self):
        """constraints.token_budget overrides config.retrieval_token_budget."""
        fts = [MockFTSResult(f"c{i}", "x" * 50, rank=-float(i)) for i in range(5)]
        vec = []
        storage = make_storage(fts_items=fts, vector_items=vec)
        tb = make_token_budget(tokens_per_call=50)
        config = MockConfig(retrieval_token_budget=10000)  # high config budget
        pipeline = make_pipeline(storage=storage, config=config, token_budget=tb)

        # Pass tight budget via constraints
        constraints = RetrievalConstraints(token_budget=99, top_k=10)
        result = await pipeline.retrieve("query", constraints=constraints)
        assert result.tokens_used <= 99

    @pytest.mark.asyncio
    async def test_memory_types_filter(self):
        """constraints.memory_types filters out non-matching item types."""
        fts = [
            MockFTSResult("claim1", "a claim", rank=-3.0),
            MockFTSResult("event1", "an event", rank=-2.0),
        ]
        # Override item_type for event
        fts[1].item_type = "event"
        vec = []
        storage = make_storage(fts_items=fts, vector_items=vec)
        pipeline = make_pipeline(storage=storage)

        constraints = RetrievalConstraints(memory_types=["claim"], top_k=10)
        result = await pipeline.retrieve("anything", constraints=constraints)
        assert all(it.item_type == "claim" for it in result.items)

    @pytest.mark.asyncio
    async def test_auto_mode_resolves_to_lightweight(self):
        """'auto' mode should resolve to 'lightweight' in Phase 1."""
        storage = make_storage()
        pipeline = make_pipeline(storage=storage)
        result = await pipeline.retrieve("query", mode="auto")
        assert result.mode == "lightweight"

    @pytest.mark.asyncio
    async def test_unknown_mode_raises_value_error(self):
        """Invalid modes raise ValueError immediately (fail-fast, #346)."""
        storage = make_storage()
        pipeline = make_pipeline(storage=storage)
        with pytest.raises(ValueError, match="Invalid retrieval mode"):
            await pipeline.retrieve("query", mode="totally_bogus_mode")

    @pytest.mark.asyncio
    async def test_agentic_mode_without_llm_falls_back_to_lightweight(self):
        """'agentic' mode without LLM still falls back to lightweight (no LLM is a valid config)."""
        storage = make_storage()
        # make_pipeline creates pipeline with no LLM by default
        pipeline = make_pipeline(storage=storage)
        result = await pipeline.retrieve("query", mode="agentic")
        assert result.mode == "lightweight"  # no-LLM fallback is expected, not a bug

    @pytest.mark.asyncio
    async def test_retrieval_id_unique_per_call(self):
        """Each retrieve() call produces a unique retrieval_id."""
        storage = make_storage()
        pipeline = make_pipeline(storage=storage)
        r1 = await pipeline.retrieve("query 1")
        r2 = await pipeline.retrieve("query 2")
        assert r1.retrieval_id != r2.retrieval_id

    @pytest.mark.asyncio
    async def test_rrf_ordering_correct(self):
        """Item appearing in both ranked lists should rank higher."""
        shared_id = "shared"
        fts = [
            MockFTSResult(shared_id, "shared content", rank=-5.0),
            MockFTSResult("fts_only", "fts only", rank=-3.0),
        ]
        vec = [
            MockVectorResult(shared_id, "shared content", score=0.9),
            MockVectorResult("vec_only", "vec only", score=0.8),
        ]
        storage = make_storage(fts_items=fts, vector_items=vec)
        pipeline = make_pipeline(storage=storage)

        result = await pipeline.retrieve("test", constraints=RetrievalConstraints(top_k=10))
        scores = {it.item_id: it.score for it in result.items}
        if shared_id in scores and "fts_only" in scores:
            assert scores[shared_id] > scores["fts_only"]
        if shared_id in scores and "vec_only" in scores:
            assert scores[shared_id] > scores["vec_only"]

    @pytest.mark.asyncio
    async def test_assemble_context_respects_budget(self):
        """assemble_context() total_tokens should not exceed budget."""
        fts = [MockFTSResult(f"c{i}", f"item {i}", rank=-float(i)) for i in range(1, 8)]
        vec = []
        storage = make_storage(fts_items=fts, vector_items=vec)
        tb = make_token_budget(tokens_per_call=20)
        config = MockConfig(retrieval_token_budget=50)  # tight budget
        pipeline = make_pipeline(storage=storage, config=config, token_budget=tb)

        ctx = await pipeline.assemble_context("test query")
        assert ctx.total_tokens <= 50
        assert ctx.budget_remaining >= 0

    @pytest.mark.asyncio
    async def test_cosine_fallback_triggered(self):
        """When vector_search raises NotImplementedError, cosine fallback is attempted."""
        fts: List = []
        storage = make_storage(fts_items=fts, vector_raises=NotImplementedError())
        # Provide get_all_embeddings for cosine fallback
        storage.get_all_embeddings = AsyncMock(return_value=[])
        pipeline = make_pipeline(storage=storage)

        # Should not raise; cosine fallback returns empty gracefully
        result = await pipeline.retrieve("query")
        assert result.mode == "lightweight"

    @pytest.mark.asyncio
    async def test_log_retrieval_fields(self):
        """Verify log_retrieval is called with all required fields."""
        fts = [MockFTSResult("m1", "content", rank=-2.0)]
        storage = make_storage(fts_items=fts, vector_items=[])
        pipeline = make_pipeline(storage=storage)

        await pipeline.retrieve("log field test", owner_id="user99")

        kwargs = storage.log_retrieval.call_args.kwargs
        required_fields = {
            "id", "owner_id", "query_text", "mode",
            "token_budget", "tokens_used", "selected_json",
            "final_context_hash", "latency_ms", "created_at",
        }
        for field_name in required_fields:
            assert field_name in kwargs, f"Missing log field: {field_name}"
        assert kwargs["owner_id"] == "user99"


# ---------------------------------------------------------------------------
# Edge case tests
# ---------------------------------------------------------------------------

class TestEdgeCases:
    @pytest.mark.asyncio
    async def test_empty_query_string(self):
        """Empty query should not crash."""
        storage = make_storage()
        pipeline = make_pipeline(storage=storage)
        result = await pipeline.retrieve("")
        assert isinstance(result, RetrievalResult)

    @pytest.mark.asyncio
    async def test_all_storages_fail(self):
        """When both FTS and vector fail, return empty results gracefully."""
        storage = make_storage(
            fts_raises=RuntimeError("DB connection lost"),
            vector_raises=RuntimeError("DB connection lost"),
        )
        pipeline = make_pipeline(storage=storage)
        result = await pipeline.retrieve("test")
        assert result.items == []
        # log_retrieval still called
        storage.log_retrieval.assert_called_once()

    @pytest.mark.asyncio
    async def test_log_failure_does_not_raise(self):
        """A failure in log_retrieval should not propagate to caller."""
        fts = [MockFTSResult("x", "content", rank=-1.0)]
        storage = make_storage(fts_items=fts)
        storage.log_retrieval = AsyncMock(side_effect=Exception("DB write error"))
        pipeline = make_pipeline(storage=storage)

        result = await pipeline.retrieve("test")
        # Should still return results despite log failure
        assert isinstance(result, RetrievalResult)

    def test_rrf_fusion_large_k(self):
        """Very large k makes all scores nearly equal."""
        items = [("a", 1.0), ("b", 0.5), ("c", 0.1)]
        result = dict(rrf_fusion([items], k=100_000))
        # All scores should be very small and close together
        scores = list(result.values())
        assert max(scores) - min(scores) < 0.0001

# ---------------------------------------------------------------------------
# Cornerstone injection tests (Issue #86)
# ---------------------------------------------------------------------------

def make_cornerstone_guardian(injection_text: str = "", raises: Optional[Exception] = None) -> MagicMock:
    """Return a mock CornerstoneGuardian with format_injection_block configured."""
    guardian = MagicMock()
    if raises is not None:
        guardian.format_injection_block = AsyncMock(side_effect=raises)
    else:
        guardian.format_injection_block = AsyncMock(return_value=injection_text)
    return guardian


class TestCornerstoneInjection:
    """Tests verifying cornerstone injection in assemble_context() (Issue #86)."""

    @pytest.mark.asyncio
    async def test_cornerstones_appear_first(self):
        """Cornerstone block must be the first part in the assembled context."""
        cs_text = "[CORNERSTONES — WHO I AM]\nname: Eva\n[END CORNERSTONES — 10 tokens]"
        fts = [MockFTSResult("m1", "retrieved memory 1", rank=-1.0)]
        storage = make_storage(fts_items=fts, vector_items=[])
        tb = make_token_budget(tokens_per_call=10)
        config = MockConfig(retrieval_token_budget=500)
        guardian = make_cornerstone_guardian(cs_text)

        pipeline = make_pipeline(storage=storage, config=config, token_budget=tb,
                                 cornerstone_guardian=guardian)
        ctx = await pipeline.assemble_context("test query")

        assert len(ctx.parts) >= 1
        first_type, first_text, _ = ctx.parts[0]
        assert first_type == "cornerstone", f"Expected 'cornerstone' first, got '{first_type}'"
        assert first_text == cs_text

    @pytest.mark.asyncio
    async def test_cornerstones_always_present_even_with_zero_retrieval(self):
        """Cornerstone block is injected even when retrieval returns 0 results."""
        cs_text = "[CORNERSTONES — WHO I AM]\nname: Eva\n[END CORNERSTONES — 5 tokens]"
        storage = make_storage(fts_items=[], vector_items=[])
        tb = make_token_budget(tokens_per_call=5)
        guardian = make_cornerstone_guardian(cs_text)

        pipeline = make_pipeline(storage=storage, token_budget=tb,
                                 cornerstone_guardian=guardian)
        ctx = await pipeline.assemble_context("empty retrieval query")

        assert len(ctx.parts) == 1
        assert ctx.parts[0][0] == "cornerstone"
        assert ctx.parts[0][1] == cs_text

    @pytest.mark.asyncio
    async def test_cornerstone_tokens_subtracted_from_retrieval_budget(self):
        """Retrieval budget passed to retrieve() must be total_budget - cornerstone_tokens."""
        cs_text = "cornerstone content"  # 10 tokens per our mock
        fts = [MockFTSResult(f"m{i}", f"memory {i}", rank=-float(i)) for i in range(1, 10)]
        storage = make_storage(fts_items=fts, vector_items=[])
        tb = make_token_budget(tokens_per_call=10)  # every call costs 10 tokens
        config = MockConfig(retrieval_token_budget=50)
        guardian = make_cornerstone_guardian(cs_text)

        pipeline = make_pipeline(storage=storage, config=config, token_budget=tb,
                                 cornerstone_guardian=guardian)
        ctx = await pipeline.assemble_context("budget test")

        # Total must not exceed budget
        assert ctx.total_tokens <= 50
        # Cornerstone is first
        assert ctx.parts[0][0] == "cornerstone"
        # Remaining parts are retrieved items (10 tokens each, 40 tokens left)
        retrieval_parts = [p for p in ctx.parts if p[0] != "cornerstone"]
        assert len(retrieval_parts) <= 4  # 40 tokens / 10 each = max 4

    @pytest.mark.asyncio
    async def test_token_budget_never_exceeded(self):
        """Total tokens in context block must never exceed retrieval_token_budget."""
        cs_text = "A cornerstone block that uses up tokens"
        fts = [MockFTSResult(f"x{i}", f"memory item {i}" * 5, rank=-float(i))
               for i in range(1, 20)]
        storage = make_storage(fts_items=fts, vector_items=[])
        tb = make_token_budget(tokens_per_call=30)
        config = MockConfig(retrieval_token_budget=100)
        guardian = make_cornerstone_guardian(cs_text)

        pipeline = make_pipeline(storage=storage, config=config, token_budget=tb,
                                 cornerstone_guardian=guardian)
        ctx = await pipeline.assemble_context("budget ceiling test")

        assert ctx.total_tokens <= 100, (
            f"Budget exceeded: {ctx.total_tokens} > 100"
        )
        assert ctx.budget_remaining >= 0

    @pytest.mark.asyncio
    async def test_no_cornerstone_guardian_still_works(self):
        """When no guardian is set, assemble_context() works as before (no cornerstone part)."""
        fts = [MockFTSResult("m1", "some memory", rank=-1.0)]
        storage = make_storage(fts_items=fts, vector_items=[])
        pipeline = make_pipeline(storage=storage)

        ctx = await pipeline.assemble_context("query without guardian")

        assert ctx.total_tokens >= 0
        assert all(p[0] != "cornerstone" for p in ctx.parts)

    @pytest.mark.asyncio
    async def test_empty_cornerstone_block_skipped(self):
        """When format_injection_block returns empty string, no cornerstone part is added."""
        storage = make_storage(fts_items=[MockFTSResult("m1", "mem", rank=-1.0)])
        guardian = make_cornerstone_guardian("")  # empty → no cornerstones sealed

        pipeline = make_pipeline(storage=storage, cornerstone_guardian=guardian)
        ctx = await pipeline.assemble_context("query")

        assert all(p[0] != "cornerstone" for p in ctx.parts)

    @pytest.mark.asyncio
    async def test_cornerstone_guardian_failure_does_not_crash(self):
        """A failure in format_injection_block should be swallowed gracefully."""
        fts = [MockFTSResult("m1", "fallback memory", rank=-1.0)]
        storage = make_storage(fts_items=fts, vector_items=[])
        guardian = make_cornerstone_guardian(raises=RuntimeError("cornerstone DB unavailable"))

        pipeline = make_pipeline(storage=storage, cornerstone_guardian=guardian)
        ctx = await pipeline.assemble_context("resilience test")

        # Should still return retrieved content without cornerstone
        assert ctx is not None
        assert all(p[0] != "cornerstone" for p in ctx.parts)

    @pytest.mark.asyncio
    async def test_call_time_guardian_overrides_instance_guardian(self):
        """A guardian passed directly to assemble_context() takes precedence."""
        instance_text = "instance cornerstone"
        call_text = "call-time cornerstone"
        fts: List = []
        storage = make_storage(fts_items=fts, vector_items=[])
        tb = make_token_budget(tokens_per_call=5)

        instance_guardian = make_cornerstone_guardian(instance_text)
        call_guardian = make_cornerstone_guardian(call_text)

        pipeline = make_pipeline(storage=storage, token_budget=tb,
                                 cornerstone_guardian=instance_guardian)
        ctx = await pipeline.assemble_context("override test", cornerstone_guardian=call_guardian)

        assert len(ctx.parts) == 1
        assert ctx.parts[0][1] == call_text
        instance_guardian.format_injection_block.assert_not_called()
        call_guardian.format_injection_block.assert_called_once()



# ---------------------------------------------------------------------------
# Sensitivity filtering tests (issue #272)
# ---------------------------------------------------------------------------

class TestSensitivityFiltering:

    def _storage_with_sens(self, items):
        fts = [MockFTSResult(i["id"], i["content"], rank=-1.0) for i in items]
        storage = make_storage(fts_items=fts, vector_items=[])
        sens = {i["id"]: i["sensitivity"] for i in items}
        storage.get_claims_sensitivity_batch = AsyncMock(return_value=sens)
        return storage

    @pytest.mark.asyncio
    async def test_restricted_excluded_when_max_secret(self):
        s = self._storage_with_sens([
            {"id": "c1", "content": "x", "sensitivity": "normal"},
            {"id": "c2", "content": "x", "sensitivity": "secret"},
            {"id": "c3", "content": "x", "sensitivity": "restricted"},
        ])
        r = await make_pipeline(storage=s).retrieve(
            "q", constraints=RetrievalConstraints(sensitivity_max="secret", top_k=10))
        ids = {i.item_id for i in r.items}
        assert "c3" not in ids and "c1" in ids and "c2" in ids

    @pytest.mark.asyncio
    async def test_secret_and_restricted_excluded_when_max_private(self):
        s = self._storage_with_sens([
            {"id": "n", "content": "x", "sensitivity": "normal"},
            {"id": "p", "content": "x", "sensitivity": "private"},
            {"id": "s", "content": "x", "sensitivity": "secret"},
            {"id": "r", "content": "x", "sensitivity": "restricted"},
        ])
        r = await make_pipeline(storage=s).retrieve(
            "q", constraints=RetrievalConstraints(sensitivity_max="private", top_k=10))
        assert {i.item_id for i in r.items} == {"n", "p"}

    @pytest.mark.asyncio
    async def test_only_normal_when_max_normal(self):
        s = self._storage_with_sens([
            {"id": "n", "content": "x", "sensitivity": "normal"},
            {"id": "p", "content": "x", "sensitivity": "private"},
        ])
        r = await make_pipeline(storage=s).retrieve(
            "q", constraints=RetrievalConstraints(sensitivity_max="normal", top_k=10))
        assert {i.item_id for i in r.items} == {"n"}

    @pytest.mark.asyncio
    async def test_claim_at_max_included(self):
        s = self._storage_with_sens([
            {"id": "a", "content": "x", "sensitivity": "private"},
            {"id": "b", "content": "x", "sensitivity": "normal"},
        ])
        r = await make_pipeline(storage=s).retrieve(
            "q", constraints=RetrievalConstraints(sensitivity_max="private", top_k=10))
        assert {i.item_id for i in r.items} == {"a", "b"}

    @pytest.mark.asyncio
    async def test_default_excludes_restricted(self):
        s = self._storage_with_sens([
            {"id": "n", "content": "x", "sensitivity": "normal"},
            {"id": "p", "content": "x", "sensitivity": "private"},
            {"id": "s", "content": "x", "sensitivity": "secret"},
            {"id": "r", "content": "x", "sensitivity": "restricted"},
        ])
        r = await make_pipeline(storage=s).retrieve("q")
        assert {i.item_id for i in r.items} == {"n", "p", "s"}

    @pytest.mark.asyncio
    async def test_missing_sensitivity_excluded(self):
        fts = [MockFTSResult("ghost", "x", rank=-1.0)]
        storage = make_storage(fts_items=fts, vector_items=[])
        storage.get_claims_sensitivity_batch = AsyncMock(return_value={})
        r = await make_pipeline(storage=storage).retrieve("q")
        assert not any(i.item_id == "ghost" for i in r.items)

    @pytest.mark.asyncio
    async def test_batch_failure_excludes_all_claims(self):
        fts = [MockFTSResult("c1", "x", rank=-1.0)]
        storage = make_storage(fts_items=fts, vector_items=[])
        storage.get_claims_sensitivity_batch = AsyncMock(side_effect=RuntimeError("db"))
        r = await make_pipeline(storage=storage).retrieve("q")
        assert not any(i.item_type == "claim" for i in r.items)

    @pytest.mark.asyncio
    async def test_events_bypass_sensitivity(self):
        fts = [MockFTSResult("ev1", "event", rank=-1.0)]
        fts[0].item_type = "event"
        storage = make_storage(fts_items=fts, vector_items=[])
        storage.get_claims_sensitivity_batch = AsyncMock(return_value={})
        r = await make_pipeline(storage=storage).retrieve(
            "q", constraints=RetrievalConstraints(sensitivity_max="normal", top_k=10))
        assert any(i.item_id == "ev1" for i in r.items)

    @pytest.mark.asyncio
    async def test_max_restricted_returns_all(self):
        s = self._storage_with_sens([
            {"id": "n", "content": "x", "sensitivity": "normal"},
            {"id": "r", "content": "x", "sensitivity": "restricted"},
        ])
        r = await make_pipeline(storage=s).retrieve(
            "q", constraints=RetrievalConstraints(sensitivity_max="restricted", top_k=10))
        assert {i.item_id for i in r.items} == {"n", "r"}


class TestSensitivityLevel:

    def test_ordering(self):
        from cortex.core.retrieval import _sensitivity_level
        assert _sensitivity_level("normal") < _sensitivity_level("private")
        assert _sensitivity_level("private") < _sensitivity_level("secret")
        assert _sensitivity_level("secret") < _sensitivity_level("restricted")

    def test_alias_sensitive_maps_to_secret(self):
        from cortex.core.retrieval import _sensitivity_level
        assert _sensitivity_level("sensitive") == _sensitivity_level("secret")

    def test_unknown_is_max(self):
        from cortex.core.retrieval import _sensitivity_level, SENSITIVITY_ORDER
        assert _sensitivity_level("bogus") > len(SENSITIVITY_ORDER) - 1


class TestQueryComplexity:
    """Tests for _is_complex_query heuristic (#495)."""

    def test_simple_queries_are_not_complex(self):
        from cortex.core.retrieval import _is_complex_query
        assert not _is_complex_query("what is my name?")
        assert not _is_complex_query("what timezone am I in?")
        assert not _is_complex_query("hello")
        assert not _is_complex_query("tell me about my hobbies")

    def test_keyword_triggers_complex(self):
        from cortex.core.retrieval import _is_complex_query
        assert _is_complex_query("what was the architecture decision for storage?")
        assert _is_complex_query("why did we choose Redis over Memcached?")
        assert _is_complex_query("compare the migration strategies")
        assert _is_complex_query("what was the tradeoff for the API design?")
        assert _is_complex_query("root cause of the outage last week")

    def test_long_query_is_complex(self):
        from cortex.core.retrieval import _is_complex_query
        long_query = "a " * 500  # 1000 chars
        assert _is_complex_query(long_query)

    def test_short_query_without_keywords_is_simple(self):
        from cortex.core.retrieval import _is_complex_query
        assert not _is_complex_query("what did I eat for lunch?")


# ---------------------------------------------------------------------------
# Status penalty tests (#574)
# ---------------------------------------------------------------------------

class TestStatusPenalty:
    """Verify that claims with non-active status get penalised in retrieval."""

    def _make_item(self, item_id: str, score: float, status: str) -> RetrievedItem:
        return RetrievedItem(
            item_type="claim",
            item_id=item_id,
            content=f"content for {item_id}",
            score=score,
            metadata={"status": status, "memory_class": "semantic", "salience": "medium"},
        )

    def test_active_claims_unpenalised(self):
        """Active claims should retain their original score."""
        pipeline = make_pipeline()
        items = [self._make_item("a", 1.0, "active")]
        result = pipeline._apply_status_penalty(items)
        assert result[0].score == 1.0

    def test_superseded_claims_heavily_penalised(self):
        """Superseded claims should have score * 0.1."""
        pipeline = make_pipeline()
        items = [self._make_item("s", 1.0, "superseded")]
        result = pipeline._apply_status_penalty(items)
        assert abs(result[0].score - 0.1) < 1e-9

    def test_deprecated_claims_penalised(self):
        """Deprecated claims should have score * 0.3."""
        pipeline = make_pipeline()
        items = [self._make_item("d", 1.0, "deprecated")]
        result = pipeline._apply_status_penalty(items)
        assert abs(result[0].score - 0.3) < 1e-9

    def test_cooling_claims_slightly_penalised(self):
        """Cooling claims should have score * 0.7."""
        pipeline = make_pipeline()
        items = [self._make_item("c", 1.0, "cooling")]
        result = pipeline._apply_status_penalty(items)
        assert abs(result[0].score - 0.7) < 1e-9

    def test_non_claim_items_unaffected(self):
        """Non-claim items (e.g. working_memory) should not be penalised."""
        pipeline = make_pipeline()
        item = RetrievedItem(
            item_type="working_memory",
            item_id="wm:1",
            content="wm content",
            score=1.0,
            metadata={"status": "superseded"},  # even with this metadata
        )
        result = pipeline._apply_status_penalty([item])
        assert result[0].score == 1.0

    def test_claims_without_metadata_unaffected(self):
        """Claims without metadata default to active (no penalty)."""
        pipeline = make_pipeline()
        item = RetrievedItem(
            item_type="claim",
            item_id="no_meta",
            content="content",
            score=1.0,
            metadata=None,
        )
        result = pipeline._apply_status_penalty([item])
        assert result[0].score == 1.0

    def test_superseded_ranks_below_active_after_penalty(self):
        """After penalty + re-sort, active claims should rank above superseded."""
        pipeline = make_pipeline()
        items = [
            self._make_item("superseded_1", 0.9, "superseded"),
            self._make_item("active_1", 0.5, "active"),
        ]
        result = pipeline._apply_status_penalty(items)
        # Superseded: 0.9 * 0.1 = 0.09, Active: 0.5 * 1.0 = 0.5
        # After sorting by score desc, active should be first
        result.sort(key=lambda it: it.score, reverse=True)
        assert result[0].item_id == "active_1"
        assert result[1].item_id == "superseded_1"

    @pytest.mark.asyncio
    async def test_status_penalty_in_full_retrieve(self):
        """End-to-end: superseded claims should score lower than active in retrieve()."""
        # Create two claims — one active, one superseded — both with same BM25 rank
        fts = [
            MockFTSResult("active_claim", "I am active", rank=-2.0),
            MockFTSResult("superseded_claim", "I am superseded", rank=-2.0),
        ]
        vec = [
            MockVectorResult("active_claim", "I am active", score=0.9),
            MockVectorResult("superseded_claim", "I am superseded", score=0.9),
        ]
        storage = make_storage(fts_items=fts, vector_items=vec)

        # Mock get_claim to return status metadata
        async def _get_claim(claim_id):
            obj = MagicMock()
            obj.subject = "test"
            obj.predicate = "is"
            obj.object_value = "something"
            obj.event_date = None
            obj.cause = None
            obj.effect = None
            obj.state_before = None
            obj.state_after = None
            obj.memory_class = "semantic"
            obj.salience = "medium"
            if claim_id == "active_claim":
                obj.status = "active"
            else:
                obj.status = "superseded"
            return obj

        storage.get_claim = AsyncMock(side_effect=_get_claim)

        pipeline = make_pipeline(storage=storage)
        result = await pipeline.retrieve("test query", constraints=RetrievalConstraints(top_k=10))

        # Find both items
        active = next((i for i in result.items if i.item_id == "active_claim"), None)
        superseded = next((i for i in result.items if i.item_id == "superseded_claim"), None)

        assert active is not None, "Active claim should be in results"
        assert superseded is not None, "Superseded claim should be in results"
        assert active.score > superseded.score, (
            f"Active claim score ({active.score}) should be higher than "
            f"superseded claim score ({superseded.score})"
        )
