"""
Tests for Voyage reranker post-retrieval integration.

Issue: #1189
"""

from __future__ import annotations

from unittest.mock import AsyncMock, patch

import pytest
import pytest_asyncio

from cortex.hippo.config import HippoConfig
from cortex.hippo.db.connection import HippoDatabase
from cortex.hippo.reranker import VoyageReranker
from cortex.hippo.retrieval import (
    HybridSearchResult,
    RerankerProvider,
    RetrievalEngine,
    pack_floats,
)
from cortex.hippo.store.conversation_store import ConversationStore
from cortex.hippo.store.summary_store import SummaryStore

# ═══════════════════════════════════════════════════════════════════════════
# Fixtures
# ═══════════════════════════════════════════════════════════════════════════

DIM = 4
OWNER = "test-owner"


def make_vec(values: list[float]) -> bytes:
    return pack_floats(values)


class MockReranker:
    """Mock reranker that reverses order."""

    async def rerank(
        self, query: str, documents: list[str], top_k: int = 10
    ) -> list[tuple[int, float]]:
        # Reverse the order and assign descending scores
        n = min(top_k, len(documents))
        return [(len(documents) - 1 - i, 1.0 - i * 0.1) for i in range(n)]


class FailingReranker:
    """Mock reranker that always raises."""

    async def rerank(
        self, query: str, documents: list[str], top_k: int = 10
    ) -> list[tuple[int, float]]:
        raise ConnectionError("API unavailable")


@pytest.fixture
def config_reranker_on() -> HippoConfig:
    return HippoConfig(
        enabled=True,
        db_path=":memory:",
        embedding_dim=DIM,
        vec_enabled=True,
        fresh_tail_count=4,
        max_tokens=10_000,
        hybrid_search_enabled=True,
        retrieval_reranker_enabled=True,
        retrieval_reranker_candidates=50,
        retrieval_reranker_top_k=5,
    )


@pytest.fixture
def config_reranker_off() -> HippoConfig:
    return HippoConfig(
        enabled=True,
        db_path=":memory:",
        embedding_dim=DIM,
        vec_enabled=True,
        fresh_tail_count=4,
        max_tokens=10_000,
        hybrid_search_enabled=True,
        retrieval_reranker_enabled=False,
    )


@pytest_asyncio.fixture
async def db_reranker(config_reranker_on) -> HippoDatabase:
    database = HippoDatabase(config_reranker_on)
    await database.initialize()
    yield database
    await database.close()


@pytest_asyncio.fixture
async def db_no_reranker(config_reranker_off) -> HippoDatabase:
    database = HippoDatabase(config_reranker_off)
    await database.initialize()
    yield database
    await database.close()


async def _create_summary(sum_store, conv_id, content, embedding_vec=None):
    summary = await sum_store.create_summary(
        conversation_id=conv_id,
        owner_id=OWNER,
        kind="leaf",
        depth=0,
        content=content,
        token_count=len(content) // 4,
    )
    if embedding_vec is not None:
        blob = make_vec(embedding_vec)
        await sum_store.update_summary_embedding(summary.summary_id, OWNER, blob)
    return summary


# ═══════════════════════════════════════════════════════════════════════════
# Tests
# ═══════════════════════════════════════════════════════════════════════════


class TestRerankerProtocol:
    """Test VoyageReranker satisfies the protocol."""

    def test_reranker_satisfies_protocol(self):
        """VoyageReranker has the required rerank method signature."""
        reranker = VoyageReranker(api_key="test-key")
        assert hasattr(reranker, "rerank")
        assert callable(reranker.rerank)

    def test_mock_reranker_satisfies_protocol(self):
        """MockReranker satisfies RerankerProvider."""
        reranker = MockReranker()
        assert hasattr(reranker, "rerank")


class TestRerankerIntegration:
    """Test reranker integration with hybrid search."""

    @pytest.mark.asyncio
    async def test_reranker_reorders_results(self, db_reranker, config_reranker_on):
        """Mock reranker changes the ordering of results."""
        conv_store = ConversationStore(db_reranker.conn)
        sum_store = SummaryStore(db_reranker.conn, vec_available=db_reranker.vec_available)
        mock_reranker = MockReranker()
        engine = RetrievalEngine(
            conversation_store=conv_store,
            summary_store=sum_store,
            db=db_reranker,
            config=config_reranker_on,
            reranker=mock_reranker,
        )
        conv = await conv_store.get_or_create_conversation(OWNER, "test-session")
        conv_id = conv.conversation_id

        s1 = await _create_summary(sum_store, conv_id, "sqlite first item", [1, 0, 0, 0])
        s2 = await _create_summary(sum_store, conv_id, "sqlite second item", [0.9, 0.1, 0, 0])
        s3 = await _create_summary(sum_store, conv_id, "sqlite third item", [0.8, 0.2, 0, 0])

        results = await engine.hybrid_search(
            query="sqlite",
            query_embedding=[1, 0, 0, 0],
            owner_id=OWNER,
            limit=5,
        )
        # MockReranker reverses the order, so results should be reordered
        assert len(results) >= 1
        assert all(isinstance(r, HybridSearchResult) for r in results)

    @pytest.mark.asyncio
    async def test_reranker_disabled_by_default(self, db_no_reranker, config_reranker_off):
        """No reranking when config is off."""
        conv_store = ConversationStore(db_no_reranker.conn)
        sum_store = SummaryStore(db_no_reranker.conn, vec_available=db_no_reranker.vec_available)
        mock_reranker = MockReranker()
        engine = RetrievalEngine(
            conversation_store=conv_store,
            summary_store=sum_store,
            db=db_no_reranker,
            config=config_reranker_off,
            reranker=mock_reranker,  # present but should NOT be used
        )
        conv = await conv_store.get_or_create_conversation(OWNER, "test-session")
        conv_id = conv.conversation_id

        await _create_summary(sum_store, conv_id, "sqlite test", [1, 0, 0, 0])

        # Wrap rerank to verify it's not called
        original_rerank = mock_reranker.rerank
        call_count = 0

        async def counting_rerank(*args, **kwargs):
            nonlocal call_count
            call_count += 1
            return await original_rerank(*args, **kwargs)

        mock_reranker.rerank = counting_rerank

        results = await engine.hybrid_search(
            query="sqlite",
            query_embedding=[1, 0, 0, 0],
            owner_id=OWNER,
        )
        assert call_count == 0, "Reranker should not be called when disabled"
        assert len(results) >= 1

    @pytest.mark.asyncio
    async def test_reranker_fallback_on_error(self, db_reranker, config_reranker_on):
        """API failure falls back to RRF scores."""
        conv_store = ConversationStore(db_reranker.conn)
        sum_store = SummaryStore(db_reranker.conn, vec_available=db_reranker.vec_available)
        engine = RetrievalEngine(
            conversation_store=conv_store,
            summary_store=sum_store,
            db=db_reranker,
            config=config_reranker_on,
            reranker=FailingReranker(),
        )
        conv = await conv_store.get_or_create_conversation(OWNER, "test-session")
        conv_id = conv.conversation_id

        await _create_summary(sum_store, conv_id, "sqlite fallback test", [1, 0, 0, 0])

        # Should not raise — falls back to RRF scores
        results = await engine.hybrid_search(
            query="sqlite",
            query_embedding=[1, 0, 0, 0],
            owner_id=OWNER,
        )
        assert len(results) >= 1

    @pytest.mark.asyncio
    async def test_reranker_top_k_respected(self, db_reranker, config_reranker_on):
        """Returns only top_k results from reranker."""
        conv_store = ConversationStore(db_reranker.conn)
        sum_store = SummaryStore(db_reranker.conn, vec_available=db_reranker.vec_available)

        class LimitedReranker:
            async def rerank(self, query, documents, top_k=10):
                n = min(top_k, len(documents))
                return [(i, 1.0 - i * 0.05) for i in range(n)]

        # Set top_k to 2 in config
        config = config_reranker_on.model_copy()
        config.retrieval_reranker_top_k = 2

        engine = RetrievalEngine(
            conversation_store=conv_store,
            summary_store=sum_store,
            db=db_reranker,
            config=config,
            reranker=LimitedReranker(),
        )
        conv = await conv_store.get_or_create_conversation(OWNER, "test-session")
        conv_id = conv.conversation_id

        for i in range(5):
            await _create_summary(sum_store, conv_id, f"sqlite item {i}", [1, 0, 0, 0])

        results = await engine.hybrid_search(
            query="sqlite",
            query_embedding=[1, 0, 0, 0],
            owner_id=OWNER,
            limit=2,
        )
        assert len(results) <= 2

    @pytest.mark.asyncio
    async def test_reranker_candidates_limit(self, db_reranker, config_reranker_on):
        """Only sends N candidates to the reranker API."""
        conv_store = ConversationStore(db_reranker.conn)
        sum_store = SummaryStore(db_reranker.conn, vec_available=db_reranker.vec_available)

        received_doc_count = []

        class CountingReranker:
            async def rerank(self, query, documents, top_k=10):
                received_doc_count.append(len(documents))
                return [(i, 1.0) for i in range(min(top_k, len(documents)))]

        config = config_reranker_on.model_copy()
        config.retrieval_reranker_candidates = 3

        engine = RetrievalEngine(
            conversation_store=conv_store,
            summary_store=sum_store,
            db=db_reranker,
            config=config,
            reranker=CountingReranker(),
        )
        conv = await conv_store.get_or_create_conversation(OWNER, "test-session")
        conv_id = conv.conversation_id

        for i in range(10):
            await _create_summary(sum_store, conv_id, f"sqlite candidate {i}", [1, 0, 0, 0])

        await engine.hybrid_search(
            query="sqlite",
            query_embedding=[1, 0, 0, 0],
            owner_id=OWNER,
        )
        assert len(received_doc_count) == 1
        assert received_doc_count[0] <= 3

    @pytest.mark.asyncio
    async def test_hybrid_with_reranker_integration(self, db_reranker, config_reranker_on):
        """End-to-end: hybrid search + rerank pipeline."""
        conv_store = ConversationStore(db_reranker.conn)
        sum_store = SummaryStore(db_reranker.conn, vec_available=db_reranker.vec_available)

        class ScoreReranker:
            """Reranker that boosts documents containing 'important'."""

            async def rerank(self, query, documents, top_k=10):
                scored = []
                for i, doc in enumerate(documents):
                    score = 0.9 if "important" in doc else 0.1
                    scored.append((i, score))
                scored.sort(key=lambda x: x[1], reverse=True)
                return scored[:top_k]

        engine = RetrievalEngine(
            conversation_store=conv_store,
            summary_store=sum_store,
            db=db_reranker,
            config=config_reranker_on,
            reranker=ScoreReranker(),
        )
        conv = await conv_store.get_or_create_conversation(OWNER, "test-session")
        conv_id = conv.conversation_id

        await _create_summary(sum_store, conv_id, "sqlite basic usage", [1, 0, 0, 0])
        s_important = await _create_summary(
            sum_store, conv_id, "sqlite important migration guide", [0.9, 0.1, 0, 0]
        )

        results = await engine.hybrid_search(
            query="sqlite",
            query_embedding=[1, 0, 0, 0],
            owner_id=OWNER,
            limit=5,
        )
        assert len(results) >= 1
        # The "important" doc should be ranked first by the reranker
        assert results[0].summary_id == s_important.summary_id

    @pytest.mark.asyncio
    async def test_reranker_zero_results(self, db_reranker, config_reranker_on):
        """Empty input handled gracefully."""
        conv_store = ConversationStore(db_reranker.conn)
        sum_store = SummaryStore(db_reranker.conn, vec_available=db_reranker.vec_available)
        engine = RetrievalEngine(
            conversation_store=conv_store,
            summary_store=sum_store,
            db=db_reranker,
            config=config_reranker_on,
            reranker=MockReranker(),
        )

        results = await engine.hybrid_search(
            query="nonexistent query xyz",
            query_embedding=[0, 0, 0, 1],
            owner_id=OWNER,
        )
        assert results == []


class TestVoyageRerankerUnit:
    """Unit tests for VoyageReranker class."""

    def test_init_defaults(self):
        """Default initialization."""
        r = VoyageReranker(api_key="test-key")
        assert r.model == "rerank-2.5"
        assert r.api_key == "test-key"
        assert "voyageai.com" in r._base_url

    def test_init_custom(self):
        """Custom model and URL."""
        r = VoyageReranker(
            api_key="key",
            model="rerank-3",
            base_url="http://localhost:8000/rerank",
        )
        assert r.model == "rerank-3"
        assert r._base_url == "http://localhost:8000/rerank"

    @pytest.mark.asyncio
    async def test_empty_documents(self):
        """Empty doc list returns empty."""
        r = VoyageReranker(api_key="test-key")
        result = await r.rerank("query", [], top_k=5)
        assert result == []
