"""
Tests for cortex/core/agentic_retrieval.py — M8 Phase 2 Agentic Retrieval.

Coverage:
  - rerank_with_feedback_and_recency: math, ordering, neutral defaults
  - deduplicate_items: keeps highest-scored copy, preserves order
  - AgenticRetrievalPipeline.retrieve(mode="agentic"):
      - cross-layer assembly (parallel layer retrieval)
      - single-round when LLM says stop
      - multi-round refinement when LLM says refine
      - max_rounds guard (never exceeds max_rounds)
      - LLM unavailable → no-LLM cross-layer fallback
      - LLM JSON parse error → graceful stop
      - token budget enforcement across rounds
      - deduplication of items appearing in multiple rounds
      - reranking with mock feedback + ages
      - log_retrieval called on every agentic retrieve
  - AgenticRetrievalPipeline._route_query: auto → agentic when LLM + complex query
  - Phase 1 backward compat: mode='lightweight' still works on AgenticRetrievalPipeline
  - RetrievalPipeline routing: mode='agentic' delegates to AgenticRetrievalPipeline
"""

from __future__ import annotations

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

import pytest

from cortex.core.agentic_retrieval import (
    AgenticRetrievalPipeline,
    AgenticRetrievalResult,
    ReasoningDecision,
    _count_by_layer,
    deduplicate_items,
    rerank_with_feedback_and_recency,
    _format_results_preview,
)
from cortex.core.retrieval import (
    RetrievalConstraints,
    RetrievalPipeline,
    RetrievalResult,
    RetrievedItem,
    _FTSResult,
    _VectorResult,
)


# ---------------------------------------------------------------------------
# Shared fixtures & helpers
# ---------------------------------------------------------------------------

@dataclass
class MockConfig:
    retrieval_token_budget: int = 2000
    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 = -2.0


@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])

    if fts_raises:
        storage.fts_search = AsyncMock(side_effect=fts_raises)
    else:
        storage.fts_search = AsyncMock(return_value=fts_items or [])

    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()
    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 = MagicMock(return_value=2000)
    tb.create_budget = MagicMock(return_value=budget_obj)
    return tb


def make_llm(response_text: Optional[str] = None, raises: Optional[Exception] = None) -> MagicMock:
    """Create a mock LLM provider.

    response_text: JSON string returned by complete(); defaults to {"should_refine": false}
    raises: if set, complete() raises this exception
    """
    llm = MagicMock()
    default_response = json.dumps({"should_refine": False, "reasoning": "Results sufficient."})
    resp = MagicMock()
    resp.text = response_text or default_response

    if raises:
        llm.complete = AsyncMock(side_effect=raises)
    else:
        llm.complete = AsyncMock(return_value=resp)

    llm.embed = AsyncMock(return_value=[0.1, 0.2, 0.3, 0.4])
    return llm


def make_pipeline(
    *,
    fts_items: Optional[List[MockFTSResult]] = None,
    vector_items: Optional[List[MockFTSResult]] = None,
    llm: Any = None,
    max_rounds: int = 3,
    config: Optional[MockConfig] = None,
    token_budget: Optional[Any] = None,
    storage: Optional[MagicMock] = None,
    feedback_manager: Any = None,
) -> AgenticRetrievalPipeline:
    if storage is None:
        storage = make_storage(fts_items=fts_items or [], vector_items=vector_items or [])
    if config is None:
        config = MockConfig()
    if token_budget is None:
        token_budget = make_token_budget()
    return AgenticRetrievalPipeline(
        storage=storage,
        config=config,
        token_budget=token_budget,
        llm=llm,
        max_rounds=max_rounds,
        feedback_manager=feedback_manager,
    )


def _item(item_id: str, score: float = 0.5, item_type: str = "claim") -> RetrievedItem:
    return RetrievedItem(item_type=item_type, item_id=item_id, content=f"content {item_id}",
                         score=score)


# ---------------------------------------------------------------------------
# Unit tests: rerank_with_feedback_and_recency
# ---------------------------------------------------------------------------

class TestRerankWithFeedbackAndRecency:
    def test_neutral_feedback_no_change_in_relative_order(self):
        """Neutral feedback (0.5) adds zero feedback contribution."""
        items = [_item("a", score=0.9), _item("b", score=0.5)]
        fb = {"a": 0.5, "b": 0.5}
        ages = {"a": 0.0, "b": 0.0}
        result = rerank_with_feedback_and_recency(items, fb, ages,
                                                  feedback_weight=0.2,
                                                  recency_weight=0.0)
        # Both get same feedback contribution — relative order preserved
        assert result[0].item_id == "a"
        assert result[1].item_id == "b"

    def test_positive_feedback_boosts_lower_scored_item(self):
        """High feedback can elevate a lower-scored item above a neutral one."""
        items = [_item("low", score=0.4), _item("high_fb", score=0.3)]
        fb = {"low": 0.5, "high_fb": 1.0}   # high_fb always helpful
        ages = {"low": 10.0, "high_fb": 10.0}
        result = rerank_with_feedback_and_recency(
            items, fb, ages,
            feedback_weight=0.5,   # large weight → feedback dominates
            recency_weight=0.0,
        )
        # high_fb: 0.3 + 0.5*(1.0-0.5)*2 = 0.3 + 0.5 = 0.8
        # low:     0.4 + 0.5*(0.5-0.5)*2 = 0.4 + 0.0 = 0.4
        assert result[0].item_id == "high_fb"

    def test_negative_feedback_penalizes_item(self):
        """Harmful feedback (ratio=0) should reduce item score."""
        items = [_item("harmful", score=0.8), _item("good", score=0.7)]
        fb = {"harmful": 0.0, "good": 0.5}
        ages = {"harmful": 0.0, "good": 0.0}
        result = rerank_with_feedback_and_recency(
            items, fb, ages,
            feedback_weight=0.3,
            recency_weight=0.0,
        )
        # harmful: 0.8 + 0.3*(0.0-0.5)*2 = 0.8 - 0.3 = 0.5
        # good:    0.7 + 0.3*(0.5-0.5)*2 = 0.7
        assert result[0].item_id == "good"

    def test_recency_bonus_exponential_decay(self):
        """Newer items should get a larger recency bonus."""
        items = [_item("old", score=0.5), _item("new", score=0.5)]
        fb = {}
        ages = {"old": 90.0, "new": 0.0}
        result = rerank_with_feedback_and_recency(
            items, fb, ages,
            feedback_weight=0.0,
            recency_weight=0.5,
            recency_half_life_days=30.0,
        )
        # new: age=0 → exp(0)=1.0 → bonus=0.5
        # old: age=90 → exp(-3*ln2) = exp(-2.08) ≈ 0.125 → bonus=0.0625
        assert result[0].item_id == "new"

    def test_scores_updated_on_returned_items(self):
        """Returned items have their .score field set to combined score."""
        items = [_item("x", score=0.6)]
        fb = {"x": 1.0}   # fully helpful
        ages = {"x": 0.0}
        result = rerank_with_feedback_and_recency(
            items, fb, ages,
            feedback_weight=0.2,
            recency_weight=0.1,
            recency_half_life_days=30.0,
        )
        # x: 0.6 + 0.2*(1.0-0.5)*2 + 0.1*exp(0) = 0.6 + 0.2 + 0.1 = 0.9
        assert abs(result[0].score - 0.9) < 1e-6

    def test_missing_feedback_defaults_neutral(self):
        """Items missing from feedback_scores use 0.5 (neutral)."""
        items = [_item("unknown", score=0.5)]
        result = rerank_with_feedback_and_recency(
            items, {}, {},
            feedback_weight=0.2,
            recency_weight=0.0,
        )
        # neutral: 0.5 + 0.2*(0.5-0.5)*2 = 0.5
        assert abs(result[0].score - 0.5) < 1e-6

    def test_empty_list(self):
        assert rerank_with_feedback_and_recency([], {}, {}) == []

    def test_result_sorted_descending(self):
        items = [_item(str(i), score=float(i) / 10) for i in range(10)]
        result = rerank_with_feedback_and_recency(items, {}, {})
        scores = [r.score for r in result]
        assert scores == sorted(scores, reverse=True)


# ---------------------------------------------------------------------------
# Unit tests: deduplicate_items
# ---------------------------------------------------------------------------

class TestDeduplicateItems:
    def test_keeps_highest_scored_duplicate(self):
        items = [_item("a", 0.9), _item("a", 0.3), _item("b", 0.7)]
        result = deduplicate_items(items)
        ids = [it.item_id for it in result]
        assert ids.count("a") == 1
        a_item = next(it for it in result if it.item_id == "a")
        assert abs(a_item.score - 0.9) < 1e-9

    def test_preserves_order_by_score(self):
        items = [_item("c", 0.5), _item("a", 0.9), _item("b", 0.7)]
        result = deduplicate_items(items)
        assert [it.item_id for it in result] == ["a", "b", "c"]

    def test_no_duplicates_unchanged(self):
        items = [_item("a", 0.9), _item("b", 0.7), _item("c", 0.5)]
        result = deduplicate_items(items)
        assert len(result) == 3

    def test_empty(self):
        assert deduplicate_items([]) == []

    def test_single_item(self):
        items = [_item("x", 0.4)]
        result = deduplicate_items(items)
        assert len(result) == 1


# ---------------------------------------------------------------------------
# Unit tests: _count_by_layer
# ---------------------------------------------------------------------------

class TestCountByLayer:
    def test_counts_correctly(self):
        items = [
            _item("a", item_type="claim"),
            _item("b", item_type="claim"),
            _item("c", item_type="event"),
            _item("d", item_type="procedure"),
        ]
        counts = _count_by_layer(items)
        assert counts == {"claim": 2, "event": 1, "procedure": 1}

    def test_empty(self):
        assert _count_by_layer([]) == {}


# ---------------------------------------------------------------------------
# Unit tests: _format_results_preview
# ---------------------------------------------------------------------------

class TestFormatResultsPreview:
    def test_no_results(self):
        preview = _format_results_preview([])
        assert preview == "(no results)"

    def test_truncates_content(self):
        item = _item("x", 0.5)
        item.content = "A" * 300
        preview = _format_results_preview([item])
        assert "A" * 200 in preview
        assert "A" * 201 not in preview

    def test_limits_items(self):
        items = [_item(str(i), 0.5) for i in range(10)]
        preview = _format_results_preview(items, max_items=3)
        # Only first 3 items listed
        assert "1." in preview
        assert "3." in preview
        assert "4." not in preview


# ---------------------------------------------------------------------------
# Integration tests: AgenticRetrievalPipeline (agentic mode)
# ---------------------------------------------------------------------------

class TestAgenticRetrievalPipeline:
    @pytest.mark.asyncio
    async def test_agentic_single_round_stop(self):
        """LLM says stop on round 1 → only 1 round executed."""
        fts = [MockFTSResult("c1", "claim content", rank=-3.0),
               MockFTSResult("c2", "more content", rank=-1.0)]
        llm = make_llm(json.dumps({"should_refine": False, "reasoning": "Good enough."}))
        pipeline = make_pipeline(fts_items=fts, llm=llm, max_rounds=3)

        result = await pipeline.retrieve("explain memory", mode="agentic")

        assert isinstance(result, AgenticRetrievalResult)
        assert result.mode == "agentic"
        # Only 1 LLM call (round 1 reasoning) → llm.complete called once
        assert llm.complete.call_count == 1
        assert result.rounds_executed == 1
        assert len(result.reasoning_trace) == 1
        assert not result.reasoning_trace[0].should_refine

    @pytest.mark.asyncio
    async def test_agentic_two_round_refinement(self):
        """LLM says refine on round 1, stop on round 2 → 2 rounds."""
        fts = [MockFTSResult("c1", "initial claim", rank=-2.0)]
        llm_responses = [
            MagicMock(text=json.dumps({
                "should_refine": True,
                "refined_query": "refined memory query",
                "reasoning": "Need more specificity.",
            })),
            MagicMock(text=json.dumps({
                "should_refine": False,
                "reasoning": "Now sufficient.",
            })),
        ]
        llm = MagicMock()
        llm.complete = AsyncMock(side_effect=llm_responses)
        llm.embed = AsyncMock(return_value=[0.1, 0.2, 0.3, 0.4])

        pipeline = make_pipeline(fts_items=fts, llm=llm, max_rounds=3)
        result = await pipeline.retrieve("explain memory", mode="agentic")

        assert isinstance(result, AgenticRetrievalResult)
        assert result.rounds_executed == 2
        assert llm.complete.call_count == 2
        assert result.reasoning_trace[0].should_refine is True
        assert result.reasoning_trace[1].should_refine is False

    @pytest.mark.asyncio
    async def test_max_rounds_guard(self):
        """LLM always says refine → execution capped at max_rounds."""
        fts = [MockFTSResult("c1", "content", rank=-2.0)]
        # Always refine with different queries to avoid seen_queries dedup
        llm_responses = [
            MagicMock(text=json.dumps({
                "should_refine": True,
                "refined_query": f"refined query {i}",
                "reasoning": f"Refine round {i}",
            }))
            for i in range(10)
        ]
        llm = MagicMock()
        llm.complete = AsyncMock(side_effect=llm_responses)
        llm.embed = AsyncMock(return_value=[0.1, 0.2, 0.3, 0.4])

        max_rounds = 3
        pipeline = make_pipeline(fts_items=fts, llm=llm, max_rounds=max_rounds)
        result = await pipeline.retrieve("find everything", mode="agentic")

        assert isinstance(result, AgenticRetrievalResult)
        # rounds_executed must not exceed max_rounds
        assert result.rounds_executed <= max_rounds
        # LLM called at most (max_rounds - 1) times (last round skips LLM)
        assert llm.complete.call_count <= max_rounds - 1

    @pytest.mark.asyncio
    async def test_no_llm_single_round_cross_layer(self):
        """Without LLM, reasoning returns stop immediately → only 1 round executed."""
        fts = [MockFTSResult("c1", "claim", rank=-2.0),
               MockFTSResult("e1", "event", rank=-1.5)]
        storage = make_storage(fts_items=fts, vector_items=[])

        pipeline = make_pipeline(storage=storage, llm=None, max_rounds=3)
        result = await pipeline.retrieve("what happened", mode="agentic")

        assert isinstance(result, AgenticRetrievalResult)
        assert result.rounds_executed == 1
        # Reasoning always returns stop when LLM unavailable;
        # trace may have 1 stop entry (round 1 reasons, finds no LLM, stops)
        if result.reasoning_trace:
            assert not result.reasoning_trace[0].should_refine
            assert "unavailable" in result.reasoning_trace[0].reasoning.lower()

    @pytest.mark.asyncio
    async def test_llm_json_parse_error_graceful_stop(self):
        """LLM returns malformed JSON → graceful stop, no crash."""
        fts = [MockFTSResult("c1", "content", rank=-2.0)]
        llm = make_llm(response_text="THIS IS NOT JSON {{{")
        pipeline = make_pipeline(fts_items=fts, llm=llm, max_rounds=3)

        result = await pipeline.retrieve("query", mode="agentic")

        assert isinstance(result, AgenticRetrievalResult)
        assert result.rounds_executed == 1
        assert result.reasoning_trace[0].should_refine is False

    @pytest.mark.asyncio
    async def test_llm_exception_graceful_stop(self):
        """LLM raises exception → graceful stop, no crash."""
        fts = [MockFTSResult("c1", "content", rank=-2.0)]
        llm = make_llm(raises=RuntimeError("API timeout"))
        pipeline = make_pipeline(fts_items=fts, llm=llm, max_rounds=3)

        result = await pipeline.retrieve("query", mode="agentic")
        assert isinstance(result, AgenticRetrievalResult)
        assert result.rounds_executed == 1

    @pytest.mark.asyncio
    async def test_deduplication_across_rounds(self):
        """Items returned in multiple rounds appear once in final result."""
        fts = [MockFTSResult("shared", "shared memory", rank=-5.0),
               MockFTSResult("unique", "unique memory", rank=-2.0)]

        # LLM triggers one refinement
        llm_responses = [
            MagicMock(text=json.dumps({
                "should_refine": True,
                "refined_query": "second query",
                "reasoning": "Refine.",
            })),
            MagicMock(text=json.dumps({"should_refine": False, "reasoning": "Done."})),
        ]
        llm = MagicMock()
        llm.complete = AsyncMock(side_effect=llm_responses)
        llm.embed = AsyncMock(return_value=[0.1, 0.2, 0.3, 0.4])

        pipeline = make_pipeline(fts_items=fts, llm=llm, max_rounds=3)
        result = await pipeline.retrieve("query", mode="agentic")

        ids = [it.item_id for it in result.items]
        # No duplicate item_ids
        assert len(ids) == len(set(ids))

    @pytest.mark.asyncio
    async def test_cross_layer_assembles_all_types(self):
        """Items from different memory types all appear in result."""
        fts_items = [
            MockFTSResult("claim1", "a claim", rank=-3.0),
            MockFTSResult("event1", "an event", rank=-2.0),
        ]
        # Override item_type for second item
        fts_items[1].item_type = "event"

        vec_items = [
            MockVectorResult("proc1", "a procedure", score=0.7, item_type="procedure"),
        ]

        storage = make_storage(fts_items=fts_items, vector_items=vec_items)
        llm = make_llm()  # stops after round 1
        pipeline = make_pipeline(storage=storage, llm=llm, max_rounds=1)

        result = await pipeline.retrieve("what is the procedure", mode="agentic")
        ids = {it.item_id for it in result.items}
        # All three items should be present (claim, event, procedure)
        assert "claim1" in ids or "event1" in ids or "proc1" in ids

    @pytest.mark.asyncio
    async def test_token_budget_enforced(self):
        """Items exceeding token budget are trimmed in agentic mode."""
        fts = [MockFTSResult(f"c{i}", f"item {i}", rank=-float(i)) for i in range(10)]
        llm = make_llm()
        tb = make_token_budget(tokens_per_call=100)
        config = MockConfig(retrieval_token_budget=250)

        pipeline = make_pipeline(
            fts_items=fts, llm=llm,
            config=config, token_budget=tb, max_rounds=1
        )
        result = await pipeline.retrieve(
            "query", mode="agentic",
            constraints=RetrievalConstraints(top_k=20)
        )
        assert result.tokens_used <= 250
        assert len(result.items) <= 3  # 250 / 100 = 2.5 → max 2 items

    @pytest.mark.asyncio
    async def test_log_retrieval_called_every_agentic_call(self):
        """log_retrieval is always called, even in agentic mode."""
        fts = [MockFTSResult("c1", "content", rank=-2.0)]
        storage = make_storage(fts_items=fts)
        llm = make_llm()
        pipeline = make_pipeline(storage=storage, llm=llm)

        await pipeline.retrieve("question", mode="agentic")
        storage.log_retrieval.assert_called_once()

        kwargs = storage.log_retrieval.call_args.kwargs
        assert kwargs["mode"] == "agentic"
        assert kwargs["query_text"] == "question"

    @pytest.mark.asyncio
    async def test_layer_counts_populated(self):
        """AgenticRetrievalResult.layer_counts reports items per memory type."""
        fts = [
            MockFTSResult("c1", "claim one", rank=-3.0),
            MockFTSResult("c2", "claim two", rank=-2.0),
        ]
        fts[1].item_type = "event"  # override second item
        llm = make_llm()
        pipeline = make_pipeline(fts_items=fts, llm=llm, max_rounds=1)

        result = await pipeline.retrieve("query", mode="agentic")
        assert isinstance(result, AgenticRetrievalResult)
        # layer_counts should be a dict
        assert isinstance(result.layer_counts, dict)

    @pytest.mark.asyncio
    async def test_reranking_applied(self):
        """Reranking is applied — mock feedback_manager called for item scores."""
        fts = [MockFTSResult("a", "content a", rank=-3.0),
               MockFTSResult("b", "content b", rank=-1.0)]
        llm = make_llm()

        # Mock FeedbackManager with get_score
        fm = MagicMock()
        feedback_score_obj = MagicMock()
        feedback_score_obj.helpfulness_ratio = 0.8
        fm.get_score = AsyncMock(return_value=feedback_score_obj)

        pipeline = make_pipeline(fts_items=fts, llm=llm, feedback_manager=fm, max_rounds=1)
        result = await pipeline.retrieve("query", mode="agentic")

        # feedback_manager.get_score should have been called for each item
        assert fm.get_score.call_count >= 1

    @pytest.mark.asyncio
    async def test_duplicate_refined_query_stops_loop(self):
        """If LLM returns the same query as before, loop terminates early."""
        fts = [MockFTSResult("c1", "content", rank=-2.0)]
        # LLM always returns same refined_query
        llm_resp = json.dumps({
            "should_refine": True,
            "refined_query": "original query",  # same as original!
            "reasoning": "Try again.",
        })
        llm = make_llm(response_text=llm_resp)
        pipeline = make_pipeline(fts_items=fts, llm=llm, max_rounds=5)

        result = await pipeline.retrieve("original query", mode="agentic")
        # Should stop after round 1 (duplicate query detected)
        assert result.rounds_executed == 1

    @pytest.mark.asyncio
    async def test_additional_filters_applied(self):
        """LLM can restrict memory_types via additional_filters."""
        # Items from multiple types
        fts = [
            MockFTSResult("c1", "a claim", rank=-3.0),
            MockFTSResult("e1", "an event", rank=-2.0),
        ]
        fts[1].item_type = "event"

        # Round 1: LLM says refine + restrict to claims only
        # Round 2: LLM says stop
        llm_responses = [
            MagicMock(text=json.dumps({
                "should_refine": True,
                "refined_query": "detailed claim query",
                "additional_filters": {"memory_types": ["claim"]},
                "reasoning": "Focus on claims.",
            })),
            MagicMock(text=json.dumps({"should_refine": False, "reasoning": "Done."})),
        ]
        llm = MagicMock()
        llm.complete = AsyncMock(side_effect=llm_responses)
        llm.embed = AsyncMock(return_value=[0.1, 0.2, 0.3, 0.4])

        storage = make_storage(fts_items=fts, vector_items=[])
        pipeline = make_pipeline(storage=storage, llm=llm, max_rounds=3)
        result = await pipeline.retrieve("find claims", mode="agentic")

        assert isinstance(result, AgenticRetrievalResult)
        assert result.rounds_executed == 2


# ---------------------------------------------------------------------------
# Integration tests: _route_query (auto mode)
# ---------------------------------------------------------------------------

class TestAgenticRouteQuery:
    def test_auto_with_llm_complex_query_routes_agentic(self):
        """Complex query + LLM available → agentic mode."""
        llm = make_llm()
        pipeline = make_pipeline(llm=llm)
        mode = pipeline._route_query("explain why the memory system works", RetrievalConstraints())
        assert mode == "agentic"

    def test_auto_with_llm_short_query_routes_lightweight(self):
        """Short, simple query → lightweight even with LLM."""
        llm = make_llm()
        pipeline = make_pipeline(llm=llm)
        mode = pipeline._route_query("memory", RetrievalConstraints())
        assert mode == "lightweight"

    def test_auto_without_llm_always_lightweight(self):
        """No LLM → always lightweight regardless of query complexity."""
        pipeline = make_pipeline(llm=None)
        mode = pipeline._route_query("explain everything about memory deeply", RetrievalConstraints())
        assert mode == "lightweight"

    @pytest.mark.asyncio
    async def test_auto_mode_uses_agentic_for_complex_query(self):
        """retrieve(mode='auto') triggers agentic path when LLM + complex query."""
        fts = [MockFTSResult("c1", "content", rank=-2.0)]
        llm = make_llm()
        pipeline = make_pipeline(fts_items=fts, llm=llm)

        result = await pipeline.retrieve(
            "explain how the memory system retrieves and ranks information", mode="auto"
        )
        assert result.mode == "agentic"

    @pytest.mark.asyncio
    async def test_auto_mode_lightweight_when_no_llm(self):
        """retrieve(mode='auto') falls back to lightweight without LLM."""
        fts = [MockFTSResult("c1", "content", rank=-2.0)]
        pipeline = make_pipeline(fts_items=fts, llm=None)
        result = await pipeline.retrieve("simple query", mode="auto")
        assert result.mode == "lightweight"


# ---------------------------------------------------------------------------
# Backward compatibility: Phase 1 behaviour on AgenticRetrievalPipeline
# ---------------------------------------------------------------------------

class TestBackwardCompatibility:
    @pytest.mark.asyncio
    async def test_lightweight_mode_still_works(self):
        """mode='lightweight' on AgenticRetrievalPipeline uses Phase 1 path."""
        fts = [MockFTSResult("c1", "content", rank=-2.0)]
        llm = make_llm()
        pipeline = make_pipeline(fts_items=fts, llm=llm)

        result = await pipeline.retrieve("query", mode="lightweight")
        assert result.mode == "lightweight"
        # LLM should NOT be called for lightweight mode
        llm.complete.assert_not_called()

    @pytest.mark.asyncio
    async def test_phase1_pipeline_delegates_agentic_to_phase2(self):
        """Phase 1 RetrievalPipeline.retrieve(mode='agentic') delegates to Phase 2 when LLM set."""
        fts = [MockFTSResult("c1", "content", rank=-2.0)]
        llm = make_llm()
        storage = make_storage(fts_items=fts)
        config = MockConfig()
        tb = make_token_budget()

        # Use base Phase 1 RetrievalPipeline (not Agentic)
        pipeline = RetrievalPipeline(
            storage=storage,
            config=config,
            token_budget=tb,
            llm=llm,
        )
        result = await pipeline.retrieve("query", mode="agentic")
        # Should produce agentic result (delegated to AgenticRetrievalPipeline)
        assert result.mode == "agentic"

    @pytest.mark.asyncio
    async def test_phase1_pipeline_falls_back_without_llm(self):
        """Phase 1 RetrievalPipeline.retrieve(mode='agentic') without LLM → lightweight."""
        fts = [MockFTSResult("c1", "content", rank=-2.0)]
        storage = make_storage(fts_items=fts)
        config = MockConfig()
        tb = make_token_budget()

        pipeline = RetrievalPipeline(storage=storage, config=config, token_budget=tb, llm=None)
        result = await pipeline.retrieve("query", mode="agentic")
        assert result.mode == "lightweight"

    @pytest.mark.asyncio
    async def test_agentic_pipeline_retrieval_id_unique(self):
        """Each agentic retrieve() produces a unique retrieval_id."""
        fts = [MockFTSResult("c1", "content", rank=-2.0)]
        llm = make_llm()
        pipeline = make_pipeline(fts_items=fts, llm=llm, max_rounds=1)

        r1 = await pipeline.retrieve("q1", mode="agentic")
        r2 = await pipeline.retrieve("q2", mode="agentic")
        assert r1.retrieval_id != r2.retrieval_id


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

class TestEdgeCases:
    @pytest.mark.asyncio
    async def test_empty_results_all_rounds(self):
        """No items from storage → empty result, no crash."""
        storage = make_storage(fts_items=[], vector_items=[])
        llm = make_llm()
        pipeline = make_pipeline(storage=storage, llm=llm, max_rounds=2)

        result = await pipeline.retrieve("query", mode="agentic")
        assert isinstance(result, AgenticRetrievalResult)
        assert result.items == []

    @pytest.mark.asyncio
    async def test_storage_layer_failure_graceful(self):
        """If one storage layer fails, others still contribute."""
        fts_items = [MockFTSResult("c1", "claim content", rank=-3.0)]
        storage = make_storage(
            fts_items=fts_items,
            vector_raises=RuntimeError("vector DB down"),
        )
        llm = make_llm()
        pipeline = make_pipeline(storage=storage, llm=llm, max_rounds=1)

        result = await pipeline.retrieve("query", mode="agentic")
        # BM25 results still present
        assert len(result.items) >= 0  # no crash

    @pytest.mark.asyncio
    async def test_max_rounds_one(self):
        """max_rounds=1 → LLM never called (only 1 round, no refinement)."""
        fts = [MockFTSResult("c1", "content", rank=-2.0)]
        llm = make_llm()
        pipeline = make_pipeline(fts_items=fts, llm=llm, max_rounds=1)

        result = await pipeline.retrieve("query", mode="agentic")
        assert result.rounds_executed == 1
        # With max_rounds=1, loop body runs once then breaks (round_num >= max_rounds)
        llm.complete.assert_not_called()

    @pytest.mark.asyncio
    async def test_assemble_context_agentic_mode(self):
        """assemble_context() works in agentic mode, respects budget."""
        fts = [MockFTSResult(f"c{i}", f"item {i}", rank=-float(i)) for i in range(5)]
        llm = make_llm()
        tb = make_token_budget(tokens_per_call=20)
        config = MockConfig(retrieval_token_budget=50)
        pipeline = make_pipeline(
            fts_items=fts, llm=llm, config=config, token_budget=tb, max_rounds=1
        )

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

    def test_rerank_with_no_feedback_manager(self):
        """Reranking works gracefully when feedback_manager is None."""
        items = [_item("a", 0.9), _item("b", 0.7)]
        result = rerank_with_feedback_and_recency(
            items, feedback_scores={}, item_ages_days={},
            feedback_weight=0.2, recency_weight=0.1
        )
        # Relative order should be preserved (both have neutral feedback)
        assert result[0].item_id == "a"
        assert result[1].item_id == "b"
