"""
Tests for Voyage rerank-2.5-lite integration in RetrievalPipeline (#502).

Coverage:
  - Rerank happy path: voyageai.AsyncClient.rerank called with correct args
  - Graceful fallback when rerank raises an exception
  - Config toggle: enabled=False skips rerank entirely
  - top_k and initial_candidates are respected
  - items are correctly mapped back to RetrievedItem objects
  - Score is updated to the rerank relevance_score
  - Empty content items don't crash rerank
  - VOYAGE_API_KEY missing triggers fallback (not an exception)
"""

from __future__ import annotations

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

import pytest

from cortex.core.retrieval import (
    RetrievalConstraints,
    RetrievalPipeline,
    RetrievedItem,
)


# ---------------------------------------------------------------------------
# Helpers & fixtures
# ---------------------------------------------------------------------------

@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"
    # Rerank config
    retrieval_rerank_enabled: bool = True
    retrieval_rerank_model: str = "rerank-2.5-lite"
    retrieval_rerank_top_k: int = 5
    retrieval_rerank_initial_candidates: int = 10


@dataclass
class MockConfigRerankDisabled(MockConfig):
    retrieval_rerank_enabled: bool = False


def make_fts_result(item_id: str, content: str, rank: float = -1.0):
    m = MagicMock()
    m.item_id = item_id
    m.content = content
    m.item_type = "claim"
    m.rank = rank
    return m


def make_vector_result(item_id: str, content: str, score: float = 0.8):
    m = MagicMock()
    m.item_id = item_id
    m.content = content
    m.item_type = "claim"
    m.score = score
    return m


def make_storage(n_items: int = 15) -> MagicMock:
    """Build a mock storage returning n_items fts + vector results."""
    storage = MagicMock()
    storage.embed = AsyncMock(return_value=[0.1, 0.2, 0.3, 0.4])

    fts_items = [make_fts_result(f"claim-{i}", f"memory content {i}", rank=-float(i + 1))
                 for i in range(n_items)]
    vec_items = [make_vector_result(f"claim-{i}", f"memory content {i}", score=1.0 - i * 0.05)
                 for i in range(n_items)]

    async def _fts(query, table="claims", top_k=50, owner_id="default"):
        return fts_items[:top_k] if table == "claims" else []

    storage.fts_search = AsyncMock(side_effect=_fts)
    storage.vector_search = AsyncMock(return_value=vec_items)
    storage.log_retrieval = AsyncMock()

    sens_map = {f"claim-{i}": "normal" for i in range(n_items)}
    storage.get_claims_sensitivity_batch = AsyncMock(return_value=sens_map)
    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 = 2000
    tb.create_budget = MagicMock(return_value=budget_obj)
    return tb


def make_pipeline(config=None, n_items: int = 15, **kwargs) -> RetrievalPipeline:
    return RetrievalPipeline(
        storage=make_storage(n_items),
        config=config or MockConfig(),
        token_budget=make_token_budget(),
        **kwargs,
    )


def make_rerank_result(indices: List[int], scores: Optional[List[float]] = None):
    """Build a fake voyageai RerankingObject."""
    if scores is None:
        scores = [1.0 - i * 0.1 for i in range(len(indices))]
    results = []
    for idx, score in zip(indices, scores):
        r = MagicMock()
        r.index = idx
        r.relevance_score = score
        results.append(r)
    obj = MagicMock()
    obj.results = results
    return obj


# ---------------------------------------------------------------------------
# _rerank_results unit tests (direct method calls)
# ---------------------------------------------------------------------------

class TestRerankResults:
    """Direct unit tests for RetrievalPipeline._rerank_results."""

    def _make_items(self, n: int) -> List[RetrievedItem]:
        return [
            RetrievedItem(
                item_type="claim",
                item_id=f"id-{i}",
                content=f"content about topic {i}",
                score=float(n - i) / n,
            )
            for i in range(n)
        ]

    @pytest.mark.asyncio
    async def test_rerank_happy_path(self):
        """Rerank reorders items by relevance_score and trims to top_k."""
        pipeline = make_pipeline()
        items = self._make_items(10)

        # Pretend reranker reverses the order (item 9 is most relevant)
        reranked_indices = list(range(9, -1, -1))
        mock_result = make_rerank_result(reranked_indices)

        mock_async_client = MagicMock()
        mock_async_client.rerank = AsyncMock(return_value=mock_result)

        with (
            patch("cortex.core.retrieval.voyageai", create=True) as mock_voyage,
            patch.dict("os.environ", {"VOYAGE_API_KEY": "test-key"}),
        ):
            mock_voyage.AsyncClient.return_value = mock_async_client
            result = await pipeline._rerank_results("test query", items)

        # Should return top_k=5 items in reranked order
        assert len(result) == 5
        # First item should be original index 9 (reversed)
        assert result[0].item_id == "id-9"
        assert result[0].score == mock_result.results[0].relevance_score
        assert result[0].provenance == "rerank:rerank-2.5-lite"

    @pytest.mark.asyncio
    async def test_rerank_called_with_correct_args(self):
        """Verify voyageai.AsyncClient.rerank is called with correct model + top_k."""
        pipeline = make_pipeline()
        items = self._make_items(8)

        mock_async_client = MagicMock()
        mock_async_client.rerank = AsyncMock(
            return_value=make_rerank_result(list(range(5)))
        )

        with (
            patch("cortex.core.retrieval.voyageai", create=True) as mock_voyage,
            patch.dict("os.environ", {"VOYAGE_API_KEY": "test-key"}),
        ):
            mock_voyage.AsyncClient.return_value = mock_async_client
            await pipeline._rerank_results("my query", items)

        mock_async_client.rerank.assert_called_once()
        call_kwargs = mock_async_client.rerank.call_args
        assert call_kwargs.kwargs.get("query") == "my query" or call_kwargs.args[0] == "my query"
        assert call_kwargs.kwargs.get("model") == "rerank-2.5-lite"
        assert call_kwargs.kwargs.get("top_k") == 5

    @pytest.mark.asyncio
    async def test_rerank_fallback_on_exception(self):
        """When rerank raises, returns original items trimmed to top_k without raising."""
        pipeline = make_pipeline()
        items = self._make_items(10)

        mock_async_client = MagicMock()
        mock_async_client.rerank = AsyncMock(side_effect=RuntimeError("API unavailable"))

        with (
            patch("cortex.core.retrieval.voyageai", create=True) as mock_voyage,
            patch.dict("os.environ", {"VOYAGE_API_KEY": "test-key"}),
        ):
            mock_voyage.AsyncClient.return_value = mock_async_client
            result = await pipeline._rerank_results("test query", items)

        # Graceful fallback: original order, trimmed to top_k
        assert len(result) == 5  # top_k
        assert result[0].item_id == "id-0"  # original order preserved

    @pytest.mark.asyncio
    async def test_rerank_fallback_when_no_api_key(self):
        """When VOYAGE_API_KEY is missing, returns original items trimmed to top_k."""
        pipeline = make_pipeline()
        items = self._make_items(10)

        with (
            patch("cortex.core.retrieval.voyageai", create=True),
            patch.dict("os.environ", {}, clear=True),
        ):
            # Ensure VOYAGE_API_KEY is absent
            import os
            os.environ.pop("VOYAGE_API_KEY", None)
            result = await pipeline._rerank_results("test query", items)

        assert len(result) == 5  # trimmed to top_k
        assert result[0].item_id == "id-0"  # original order

    @pytest.mark.asyncio
    async def test_rerank_preserves_item_fields(self):
        """Reranked items carry original item_type, item_id, content."""
        pipeline = make_pipeline()
        items = [
            RetrievedItem(item_type="event", item_id="ev-1", content="an event", score=0.5),
            RetrievedItem(item_type="claim", item_id="cl-1", content="a claim", score=0.4),
        ]
        mock_result = make_rerank_result([1, 0], scores=[0.95, 0.80])

        mock_async_client = MagicMock()
        mock_async_client.rerank = AsyncMock(return_value=mock_result)

        with (
            patch("cortex.core.retrieval.voyageai", create=True) as mock_voyage,
            patch.dict("os.environ", {"VOYAGE_API_KEY": "test-key"}),
        ):
            mock_voyage.AsyncClient.return_value = mock_async_client
            result = await pipeline._rerank_results("query", items)

        assert result[0].item_id == "cl-1"
        assert result[0].item_type == "claim"
        assert result[0].content == "a claim"
        assert result[0].score == 0.95

    @pytest.mark.asyncio
    async def test_rerank_empty_content_items(self):
        """Items with all-empty content fall back gracefully (no API call crash)."""
        pipeline = make_pipeline()
        items = [
            RetrievedItem(item_type="claim", item_id="id-0", content="   ", score=0.9),
            RetrievedItem(item_type="claim", item_id="id-1", content="", score=0.8),
        ]

        with (
            patch("cortex.core.retrieval.voyageai", create=True),
            patch.dict("os.environ", {"VOYAGE_API_KEY": "test-key"}),
        ):
            # Should not raise, should return trimmed items
            result = await pipeline._rerank_results("query", items)

        assert isinstance(result, list)
        assert len(result) <= 5


# ---------------------------------------------------------------------------
# Config toggle tests
# ---------------------------------------------------------------------------

class TestRerankConfigToggle:
    """Verify rerank is skipped when disabled in config."""

    @pytest.mark.asyncio
    async def test_rerank_disabled_skips_voyage_call(self):
        """When retrieval_rerank_enabled=False, voyageai is never imported/called."""
        pipeline = make_pipeline(config=MockConfigRerankDisabled(), n_items=5)

        with patch("cortex.core.retrieval.voyageai", create=True) as mock_voyage:
            with patch.dict("os.environ", {"VOYAGE_API_KEY": "test-key"}):
                result = await pipeline.retrieve("what do I know about topic 0")

        # voyageai.AsyncClient should not be instantiated
        mock_voyage.AsyncClient.assert_not_called()
        assert isinstance(result.items, list)

    @pytest.mark.asyncio
    async def test_rerank_enabled_calls_voyage(self):
        """When retrieval_rerank_enabled=True, voyageai.AsyncClient is used."""
        pipeline = make_pipeline(config=MockConfig(), n_items=5)

        mock_async_client = MagicMock()
        mock_async_client.rerank = AsyncMock(
            return_value=make_rerank_result(list(range(min(5, 5))))
        )

        with (
            patch("cortex.core.retrieval.voyageai", create=True) as mock_voyage,
            patch.dict("os.environ", {"VOYAGE_API_KEY": "test-key"}),
        ):
            mock_voyage.AsyncClient.return_value = mock_async_client
            await pipeline.retrieve("test query")

        mock_voyage.AsyncClient.assert_called_once_with(api_key="test-key")
        mock_async_client.rerank.assert_called_once()


# ---------------------------------------------------------------------------
# Initial candidates test
# ---------------------------------------------------------------------------

class TestRerankInitialCandidates:
    """Verify initial_candidates controls first-pass retrieval width."""

    @pytest.mark.asyncio
    async def test_initial_candidates_expands_top_k(self):
        """When initial_candidates > constraints.top_k, retrieval fetches more."""

        @dataclass
        class WideConfig(MockConfig):
            retrieval_rerank_initial_candidates: int = 10
            retrieval_rerank_top_k: int = 3

        captured_top_k = []

        storage = make_storage(n_items=12)

        async def _fts_capture(query, table="claims", top_k=50, owner_id="default"):
            if table == "claims":
                captured_top_k.append(top_k)
            return [make_fts_result(f"claim-{i}", f"content {i}", -float(i + 1))
                    for i in range(min(top_k, 12))]

        storage.fts_search = AsyncMock(side_effect=_fts_capture)
        storage.vector_search = AsyncMock(return_value=[
            make_vector_result(f"claim-{i}", f"content {i}", 0.9 - i * 0.05)
            for i in range(12)
        ])
        sens_map = {f"claim-{i}": "normal" for i in range(12)}
        storage.get_claims_sensitivity_batch = AsyncMock(return_value=sens_map)

        pipeline = RetrievalPipeline(storage, WideConfig(), make_token_budget())

        mock_async_client = MagicMock()
        mock_async_client.rerank = AsyncMock(
            return_value=make_rerank_result(list(range(3)))
        )

        with (
            patch("cortex.core.retrieval.voyageai", create=True) as mock_voyage,
            patch.dict("os.environ", {"VOYAGE_API_KEY": "test-key"}),
        ):
            mock_voyage.AsyncClient.return_value = mock_async_client
            # Default constraints.top_k = 20, but initial_candidates = 10
            # Since initial_candidates (10) < constraints.top_k (20), no override
            await pipeline.retrieve("test query", constraints=RetrievalConstraints(top_k=5))

        # FTS was called with max(5, 10)=10 * candidate_multiplier=3 = 30
        assert any(k >= 30 for k in captured_top_k), (
            f"Expected top_k >= 30 in FTS calls, got {captured_top_k}"
        )

    @pytest.mark.asyncio
    async def test_final_results_capped_at_rerank_top_k(self):
        """Result list after rerank is capped at rerank_top_k (5)."""
        pipeline = make_pipeline(config=MockConfig(), n_items=12)

        # Return 5 reranked results
        mock_async_client = MagicMock()
        mock_async_client.rerank = AsyncMock(
            return_value=make_rerank_result(list(range(5)), scores=[0.9, 0.8, 0.7, 0.6, 0.5])
        )

        with (
            patch("cortex.core.retrieval.voyageai", create=True) as mock_voyage,
            patch.dict("os.environ", {"VOYAGE_API_KEY": "test-key"}),
        ):
            mock_voyage.AsyncClient.return_value = mock_async_client
            result = await pipeline.retrieve("test query")

        # items should be <= top_k (5) modulo token budget
        assert len(result.items) <= 5


# ---------------------------------------------------------------------------
# End-to-end integration: retrieve() with rerank
# ---------------------------------------------------------------------------

class TestRetrieveWithRerank:
    """End-to-end retrieve() tests with mocked voyageai."""

    @pytest.mark.asyncio
    async def test_retrieve_returns_reranked_order(self):
        """retrieve() returns items in the order provided by Voyage reranker."""
        pipeline = make_pipeline(config=MockConfig(), n_items=8)

        # Reranker says item 7 is most relevant, then 5, 3, 1, 0
        reranked_indices = [7, 5, 3, 1, 0]
        mock_async_client = MagicMock()
        mock_async_client.rerank = AsyncMock(
            return_value=make_rerank_result(reranked_indices, scores=[0.99, 0.88, 0.77, 0.66, 0.55])
        )

        with (
            patch("cortex.core.retrieval.voyageai", create=True) as mock_voyage,
            patch.dict("os.environ", {"VOYAGE_API_KEY": "test-key"}),
        ):
            mock_voyage.AsyncClient.return_value = mock_async_client
            result = await pipeline.retrieve("test query")

        # First item should have score from reranker
        if result.items:
            assert result.items[0].score == 0.99
            assert result.items[0].provenance == "rerank:rerank-2.5-lite"

    @pytest.mark.asyncio
    async def test_retrieve_fallback_on_rerank_failure(self):
        """retrieve() still returns results when rerank raises, just unranked."""
        pipeline = make_pipeline(config=MockConfig(), n_items=5)

        mock_async_client = MagicMock()
        mock_async_client.rerank = AsyncMock(side_effect=ConnectionError("Network error"))

        with (
            patch("cortex.core.retrieval.voyageai", create=True) as mock_voyage,
            patch.dict("os.environ", {"VOYAGE_API_KEY": "test-key"}),
        ):
            mock_voyage.AsyncClient.return_value = mock_async_client
            result = await pipeline.retrieve("test query")

        # Must still return a valid RetrievalResult
        assert result is not None
        assert hasattr(result, "items")
        assert isinstance(result.items, list)
        # No provenance set (fallback doesn't add rerank provenance)
        for item in result.items:
            assert item.provenance != "rerank:rerank-2.5-lite"

    @pytest.mark.asyncio
    async def test_retrieve_api_interface_unchanged(self):
        """retrieve() signature and return type are unchanged by rerank addition."""
        pipeline = make_pipeline(config=MockConfig(), n_items=3)

        mock_async_client = MagicMock()
        mock_async_client.rerank = AsyncMock(
            return_value=make_rerank_result([0, 1, 2])
        )

        with (
            patch("cortex.core.retrieval.voyageai", create=True) as mock_voyage,
            patch.dict("os.environ", {"VOYAGE_API_KEY": "test-key"}),
        ):
            mock_voyage.AsyncClient.return_value = mock_async_client
            result = await pipeline.retrieve(
                "test",
                mode="auto",
                constraints=RetrievalConstraints(top_k=3),
                owner_id="default",
            )

        # RetrievalResult interface intact
        assert hasattr(result, "items")
        assert hasattr(result, "retrieval_id")
        assert hasattr(result, "mode")
        assert hasattr(result, "latency_ms")
        assert hasattr(result, "tokens_used")
