"""Tests for chained multi-hop retrieval (#1279)."""
from __future__ import annotations

import pytest
from unittest.mock import AsyncMock, MagicMock
from dataclasses import dataclass

from cortex.core.agentic_retrieval import (
    extract_entities_from_items,
    chained_retrieve,
)
from cortex.core.retrieval_types import RetrievedItem, RetrievalResult


class TestExtractEntities:
    def test_extracts_from_metadata(self):
        items = [
            RetrievedItem(
                item_type="claim", item_id="1",
                content="User likes hiking",
                score=0.8,
                metadata={"subject": "Andrew"},
            ),
        ]
        entities = extract_entities_from_items(items)
        assert "Andrew" in entities

    def test_extracts_proper_nouns_from_content(self):
        items = [
            RetrievedItem(
                item_type="claim", item_id="1",
                content="Andrew moved to Bangkok last year",
                score=0.8,
            ),
        ]
        entities = extract_entities_from_items(items)
        assert "Andrew" in entities
        assert "Bangkok" in entities

    def test_filters_common_words(self):
        items = [
            RetrievedItem(
                item_type="claim", item_id="1",
                content="The user has been working There since January",
                score=0.8,
            ),
        ]
        entities = extract_entities_from_items(items)
        # "The" and "There" should be filtered
        assert "The" not in entities
        assert "There" not in entities

    def test_empty_items(self):
        assert extract_entities_from_items([]) == []

    def test_caps_at_20(self):
        items = [
            RetrievedItem(
                item_type="claim", item_id=str(i),
                content=f"Person{i} lives in City{i}",
                score=0.5,
                metadata={"subject": f"Entity{i}"},
            )
            for i in range(30)
        ]
        entities = extract_entities_from_items(items)
        assert len(entities) <= 20


class TestChainedRetrieve:
    @pytest.mark.asyncio
    async def test_basic_chaining(self):
        """Chained retrieval makes multiple hops."""
        call_count = 0

        class MockPipeline:
            async def retrieve(self, query, owner_id="default", constraints=None, mode="lightweight"):
                nonlocal call_count
                call_count += 1
                if call_count == 1:
                    return RetrievalResult(
                        items=[
                            RetrievedItem(
                                item_type="claim", item_id="c1",
                                content="Andrew lives in Bangkok",
                                score=0.9,
                                metadata={"subject": "Andrew"},
                            ),
                        ],
                        retrieval_id="r1", mode="lightweight",
                        latency_ms=10, tokens_used=5,
                    )
                else:
                    return RetrievalResult(
                        items=[
                            RetrievedItem(
                                item_type="claim", item_id="c2",
                                content="Bangkok is in Thailand",
                                score=0.7,
                            ),
                        ],
                        retrieval_id="r2", mode="lightweight",
                        latency_ms=10, tokens_used=5,
                    )

        result = await chained_retrieve(
            MockPipeline(), "Where does Andrew live?", max_hops=2,
        )
        assert len(result.items) == 2
        assert call_count >= 2
        assert "chained" in result.mode

    @pytest.mark.asyncio
    async def test_stops_when_no_new_items(self):
        """Chaining stops when follow-up queries don't return new items."""
        class MockPipeline:
            async def retrieve(self, query, owner_id="default", constraints=None, mode="lightweight"):
                return RetrievalResult(
                    items=[
                        RetrievedItem(
                            item_type="claim", item_id="c1",
                            content="Static result", score=0.5,
                        ),
                    ],
                    retrieval_id="r1", mode="lightweight",
                    latency_ms=10, tokens_used=5,
                )

        result = await chained_retrieve(
            MockPipeline(), "test query", max_hops=3,
        )
        # Should stop early since no new items discovered
        assert len(result.items) == 1

    @pytest.mark.asyncio
    async def test_deduplicates_across_hops(self):
        """Items seen in earlier hops aren't duplicated."""
        class MockPipeline:
            async def retrieve(self, query, owner_id="default", constraints=None, mode="lightweight"):
                return RetrievalResult(
                    items=[
                        RetrievedItem(
                            item_type="claim", item_id="c1",
                            content="Same claim", score=0.5,
                        ),
                        RetrievedItem(
                            item_type="claim", item_id="c2",
                            content="Another claim about Person1",
                            score=0.4,
                            metadata={"subject": "Person1"},
                        ),
                    ],
                    retrieval_id="r1", mode="lightweight",
                    latency_ms=10, tokens_used=5,
                )

        result = await chained_retrieve(
            MockPipeline(), "test query", max_hops=2,
        )
        ids = [it.item_id for it in result.items]
        assert len(ids) == len(set(ids))  # no duplicates

    @pytest.mark.asyncio
    async def test_keeps_highest_scored_duplicate(self):
        """When a later hop returns the same item with a higher score, keep the best."""
        call_count = 0

        class MockPipeline:
            async def retrieve(self, query, owner_id="default", constraints=None, mode="lightweight"):
                nonlocal call_count
                call_count += 1
                if call_count == 1:
                    return RetrievalResult(
                        items=[
                            RetrievedItem(
                                item_type="claim", item_id="c1",
                                content="Andrew lives in Bangkok",
                                score=0.5,
                                metadata={"subject": "Andrew"},
                            ),
                        ],
                        retrieval_id="r1", mode="lightweight",
                        latency_ms=10, tokens_used=5,
                    )
                else:
                    return RetrievalResult(
                        items=[
                            RetrievedItem(
                                item_type="claim", item_id="c1",
                                content="Andrew lives in Bangkok",
                                score=0.9,  # higher score in later hop
                            ),
                            RetrievedItem(
                                item_type="claim", item_id="c2",
                                content="Bangkok is in Thailand",
                                score=0.7,
                            ),
                        ],
                        retrieval_id="r2", mode="lightweight",
                        latency_ms=10, tokens_used=5,
                    )

        result = await chained_retrieve(
            MockPipeline(), "Where does Andrew live?", max_hops=2,
        )
        # c1 should have the higher score (0.9) from the second hop
        c1 = [it for it in result.items if it.item_id == "c1"][0]
        assert c1.score == 0.9
