"""Tests for entity_keys filter in retrieve() and get_all()."""

from __future__ import annotations

import uuid
from contextlib import asynccontextmanager
from datetime import UTC, datetime
from unittest.mock import AsyncMock, MagicMock, patch

import pytest

from arandu.client import MemoryClient

AGENT_ID = "test-agent"
NOW = datetime.now(tz=UTC)


def _make_client() -> MemoryClient:
    """Create a MemoryClient with mocked providers (no real DB)."""
    with patch("arandu.client.create_engine"), patch("arandu.client.create_session_factory"):
        return MemoryClient(
            database_url="postgresql+psycopg://fake:fake@localhost/fake",
            llm=MagicMock(),
            embeddings=MagicMock(),
        )


def _mock_session_factory(session: AsyncMock) -> MagicMock:
    @asynccontextmanager
    async def _factory():  # type: ignore[no-untyped-def]
        yield session

    return _factory


def _make_orm_fact(
    fact_id: uuid.UUID | None = None,
    entity_key: str = "person:pedro",
    entity_name: str = "Pedro",
    fact_text: str = "Pedro lives in São Paulo",
    **overrides: object,
) -> MagicMock:
    defaults = {
        "id": fact_id or uuid.uuid4(),
        "agent_id": AGENT_ID,
        "entity_name": entity_name,
        "entity_key": entity_key,
        "entity_type": "person",
        "attribute_key": None,
        "fact_text": fact_text,
        "category": None,
        "confidence": 0.9,
        "importance": 0.5,
        "valid_from": NOW,
        "created_at": NOW,
        "source_context": None,
        "speaker": "Pedro",
    }
    defaults.update(overrides)
    mock = MagicMock()
    for k, v in defaults.items():
        setattr(mock, k, v)
    return mock


class TestGetAllWithEntityKeys:
    """Tests for get_all() with entity_keys filter."""

    @pytest.mark.asyncio
    async def test_entity_keys_filter_calls_execute(self) -> None:
        """Verifies the query is executed when entity_keys is provided."""
        facts = [_make_orm_fact()]

        session = AsyncMock()
        mock_result = MagicMock()
        mock_result.scalars.return_value.all.return_value = facts
        session.execute.return_value = mock_result

        client = _make_client()
        client._session_factory = _mock_session_factory(session)

        # Patch resolver so this test focuses on get_all filter wiring only.
        with patch(
            "arandu.read.retrieval._resolve_entity_keys",
            new_callable=AsyncMock,
            return_value=(["person:pedro"], []),
        ):
            details = await client.get_all(AGENT_ID, entity_keys=["person:pedro"])

        assert len(details) == 1
        assert details[0].entity_key == "person:pedro"
        session.execute.assert_called_once()

    @pytest.mark.asyncio
    async def test_entity_keys_none_is_backward_compat(self) -> None:
        """Without entity_keys, behavior is identical to before."""
        facts = [_make_orm_fact(), _make_orm_fact(entity_key="person:ana", entity_name="Ana")]

        session = AsyncMock()
        mock_result = MagicMock()
        mock_result.scalars.return_value.all.return_value = facts
        session.execute.return_value = mock_result

        client = _make_client()
        client._session_factory = _mock_session_factory(session)

        details = await client.get_all(AGENT_ID)

        assert len(details) == 2

    @pytest.mark.asyncio
    async def test_entity_keys_empty_result(self) -> None:
        """Entity key that matches nothing returns empty list."""
        session = AsyncMock()
        mock_result = MagicMock()
        mock_result.scalars.return_value.all.return_value = []
        session.execute.return_value = mock_result

        client = _make_client()
        client._session_factory = _mock_session_factory(session)

        with patch(
            "arandu.read.retrieval._resolve_entity_keys",
            new_callable=AsyncMock,
            return_value=(["person:nonexistent"], []),
        ):
            details = await client.get_all(AGENT_ID, entity_keys=["person:nonexistent"])

        assert details == []

    @pytest.mark.asyncio
    async def test_multiple_entity_keys(self) -> None:
        """Multiple entity_keys work as OR — returns facts for any of them."""
        facts = [
            _make_orm_fact(entity_key="person:pedro"),
            _make_orm_fact(entity_key="person:ana", entity_name="Ana"),
        ]

        session = AsyncMock()
        mock_result = MagicMock()
        mock_result.scalars.return_value.all.return_value = facts
        session.execute.return_value = mock_result

        client = _make_client()
        client._session_factory = _mock_session_factory(session)

        with patch(
            "arandu.read.retrieval._resolve_entity_keys",
            new_callable=AsyncMock,
            return_value=(["person:pedro", "person:ana"], []),
        ):
            details = await client.get_all(AGENT_ID, entity_keys=["person:pedro", "person:ana"])

        assert len(details) == 2


class TestRetrieveWithEntityKeys:
    """Tests for retrieve() with entity_keys parameter."""

    @pytest.mark.asyncio
    async def test_entity_keys_propagated_to_pipeline(self) -> None:
        """Verifies entity_keys is passed through to run_read_pipeline."""
        client = _make_client()

        mock_read_result = MagicMock()
        mock_read_result.facts = []
        mock_read_result.context = ""
        mock_read_result.total_candidates = 0
        mock_read_result.duration_ms = 0.0

        session = AsyncMock()

        client._session_factory = _mock_session_factory(session)

        with patch("arandu.read.pipeline.run_read_pipeline", new_callable=AsyncMock) as mock_pipeline:
            mock_pipeline.return_value = mock_read_result

            await client.retrieve(AGENT_ID, "query", entity_keys=["person:pedro"])

            mock_pipeline.assert_called_once()
            assert mock_pipeline.call_args[1]["entity_keys"] == ["person:pedro"]

    @pytest.mark.asyncio
    async def test_entity_keys_none_by_default(self) -> None:
        """Without entity_keys, None is passed to the pipeline."""
        client = _make_client()

        mock_read_result = MagicMock()
        mock_read_result.facts = []
        mock_read_result.context = ""
        mock_read_result.total_candidates = 0
        mock_read_result.duration_ms = 0.0

        session = AsyncMock()
        client._session_factory = _mock_session_factory(session)

        with patch("arandu.read.pipeline.run_read_pipeline", new_callable=AsyncMock) as mock_pipeline:
            mock_pipeline.return_value = mock_read_result

            await client.retrieve(AGENT_ID, "query")

            assert mock_pipeline.call_args[1]["entity_keys"] is None


# ---------------------------------------------------------------------------
# Alias resolution tests
# ---------------------------------------------------------------------------


def _make_alias_row(alias: str, canonical_key: str) -> tuple[str, str]:
    """Build a row as returned by the alias lookup query."""
    return (alias, canonical_key)


def _make_entity_row(canonical_key: str) -> tuple[str]:
    """Build a row as returned by the canonical lookup query."""
    return (canonical_key,)


class TestResolveEntityKeys:
    """Tests for the ``_resolve_entity_keys`` helper."""

    @pytest.mark.asyncio
    async def test_canonical_key_passes_through(self) -> None:
        """A key that matches a canonical entity is returned unchanged."""
        from arandu.read.retrieval import _resolve_entity_keys

        session = AsyncMock()

        canonical_result = MagicMock()
        canonical_result.all.return_value = [_make_entity_row("person:pedro_menezes")]
        # Second query (aliases) should not be needed because everything matched
        alias_result = MagicMock()
        alias_result.all.return_value = []
        session.execute = AsyncMock(side_effect=[canonical_result, alias_result])

        resolved, invalid = await _resolve_entity_keys(session, AGENT_ID, ["person:pedro_menezes"])

        assert resolved == ["person:pedro_menezes"]
        assert invalid == []

    @pytest.mark.asyncio
    async def test_alias_resolves_to_canonical(self) -> None:
        """An alias is replaced with its canonical key."""
        from arandu.read.retrieval import _resolve_entity_keys

        session = AsyncMock()

        # Canonical lookup finds nothing — "person:pedro" is not canonical
        canonical_result = MagicMock()
        canonical_result.all.return_value = []
        # Alias lookup finds "pedro" → "person:pedro_menezes"
        alias_result = MagicMock()
        alias_result.all.return_value = [_make_alias_row("pedro", "person:pedro_menezes")]
        session.execute = AsyncMock(side_effect=[canonical_result, alias_result])

        resolved, invalid = await _resolve_entity_keys(session, AGENT_ID, ["person:pedro"])

        assert resolved == ["person:pedro_menezes"]
        assert invalid == []

    @pytest.mark.asyncio
    async def test_alias_with_full_prefix_form(self) -> None:
        """An alias that matches the full key form (``person:pedro``) resolves."""
        from arandu.read.retrieval import _resolve_entity_keys

        session = AsyncMock()

        canonical_result = MagicMock()
        canonical_result.all.return_value = []
        alias_result = MagicMock()
        alias_result.all.return_value = [_make_alias_row("person:pedro", "person:pedro_menezes")]
        session.execute = AsyncMock(side_effect=[canonical_result, alias_result])

        resolved, invalid = await _resolve_entity_keys(session, AGENT_ID, ["person:pedro"])

        assert resolved == ["person:pedro_menezes"]
        assert invalid == []

    @pytest.mark.asyncio
    async def test_unknown_key_goes_to_invalid(self) -> None:
        """A key that matches neither canonical nor alias is marked invalid."""
        from arandu.read.retrieval import _resolve_entity_keys

        session = AsyncMock()

        canonical_result = MagicMock()
        canonical_result.all.return_value = []
        alias_result = MagicMock()
        alias_result.all.return_value = []
        session.execute = AsyncMock(side_effect=[canonical_result, alias_result])

        resolved, invalid = await _resolve_entity_keys(session, AGENT_ID, ["person:unknown"])

        assert resolved == []
        assert invalid == ["person:unknown"]

    @pytest.mark.asyncio
    async def test_partial_mix_valid_and_invalid(self) -> None:
        """A mix of valid + invalid returns both lists correctly populated."""
        from arandu.read.retrieval import _resolve_entity_keys

        session = AsyncMock()

        canonical_result = MagicMock()
        canonical_result.all.return_value = [_make_entity_row("person:pedro_menezes")]
        alias_result = MagicMock()
        alias_result.all.return_value = []
        session.execute = AsyncMock(side_effect=[canonical_result, alias_result])

        resolved, invalid = await _resolve_entity_keys(session, AGENT_ID, ["person:pedro_menezes", "person:xyz"])

        assert resolved == ["person:pedro_menezes"]
        assert invalid == ["person:xyz"]

    @pytest.mark.asyncio
    async def test_empty_list_returns_empty(self) -> None:
        """Empty input returns two empty lists without querying the DB."""
        from arandu.read.retrieval import _resolve_entity_keys

        session = AsyncMock()
        session.execute = AsyncMock()

        resolved, invalid = await _resolve_entity_keys(session, AGENT_ID, [])

        assert resolved == []
        assert invalid == []
        session.execute.assert_not_called()

    @pytest.mark.asyncio
    async def test_failsafe_on_db_error(self) -> None:
        """When the DB query raises, fall back to the raw list with no invalids."""
        from arandu.read.retrieval import _resolve_entity_keys

        session = AsyncMock()
        session.execute = AsyncMock(side_effect=RuntimeError("connection lost"))

        resolved, invalid = await _resolve_entity_keys(session, AGENT_ID, ["person:pedro", "person:ana"])

        assert resolved == ["person:pedro", "person:ana"]
        assert invalid == []

    @pytest.mark.asyncio
    async def test_duplicates_are_deduplicated(self) -> None:
        """When two different keys resolve to the same canonical, dedup applies."""
        from arandu.read.retrieval import _resolve_entity_keys

        session = AsyncMock()

        canonical_result = MagicMock()
        canonical_result.all.return_value = []
        alias_result = MagicMock()
        alias_result.all.return_value = [
            _make_alias_row("pedro", "person:pedro_menezes"),
            _make_alias_row("pedro menezes", "person:pedro_menezes"),
        ]
        session.execute = AsyncMock(side_effect=[canonical_result, alias_result])

        resolved, invalid = await _resolve_entity_keys(session, AGENT_ID, ["person:pedro", "person:pedro menezes"])

        assert resolved == ["person:pedro_menezes"]
        assert invalid == []


# ---------------------------------------------------------------------------
# retrieve() warnings propagation
# ---------------------------------------------------------------------------


class TestRetrieveWarnings:
    """Tests that warnings flow from ReadResult into RetrieveResult."""

    @pytest.mark.asyncio
    async def test_warnings_propagated_from_pipeline(self) -> None:
        """Warnings on the internal ReadResult appear on the public RetrieveResult."""
        client = _make_client()

        mock_read_result = MagicMock()
        mock_read_result.facts = []
        mock_read_result.context = ""
        mock_read_result.total_candidates = 0
        mock_read_result.duration_ms = 0.0
        mock_read_result.warnings = ["entity_key 'person:xyz' not found (not canonical, no matching alias)"]

        session = AsyncMock()
        client._session_factory = _mock_session_factory(session)

        with patch(
            "arandu.read.pipeline.run_read_pipeline",
            new_callable=AsyncMock,
            return_value=mock_read_result,
        ):
            result = await client.retrieve(AGENT_ID, "query", entity_keys=["person:xyz"])

            assert len(result.warnings) == 1
            assert "person:xyz" in result.warnings[0]
            assert "not found" in result.warnings[0]

    @pytest.mark.asyncio
    async def test_warnings_empty_when_all_resolved(self) -> None:
        """When every key resolves, warnings is an empty list."""
        client = _make_client()

        mock_read_result = MagicMock()
        mock_read_result.facts = []
        mock_read_result.context = ""
        mock_read_result.total_candidates = 0
        mock_read_result.duration_ms = 0.0
        mock_read_result.warnings = []

        session = AsyncMock()
        client._session_factory = _mock_session_factory(session)

        with patch(
            "arandu.read.pipeline.run_read_pipeline",
            new_callable=AsyncMock,
            return_value=mock_read_result,
        ):
            result = await client.retrieve(AGENT_ID, "query", entity_keys=["person:pedro_menezes"])

            assert result.warnings == []


# ---------------------------------------------------------------------------
# get_all() alias resolution
# ---------------------------------------------------------------------------


class TestGetAllAliasResolution:
    """Tests that get_all() resolves aliases before filtering."""

    @pytest.mark.asyncio
    async def test_get_all_resolves_alias_before_query(self) -> None:
        """get_all() calls _resolve_entity_keys before building the filter."""
        client = _make_client()

        session = AsyncMock()
        mock_result = MagicMock()
        mock_result.scalars.return_value.all.return_value = []
        session.execute.return_value = mock_result
        client._session_factory = _mock_session_factory(session)

        with patch(
            "arandu.read.retrieval._resolve_entity_keys",
            new_callable=AsyncMock,
            return_value=(["person:pedro_menezes"], []),
        ) as mock_resolve:
            await client.get_all(AGENT_ID, entity_keys=["person:pedro"])

            mock_resolve.assert_called_once()
            call_args = mock_resolve.call_args[0]
            assert call_args[2] == ["person:pedro"]

    @pytest.mark.asyncio
    async def test_get_all_without_entity_keys_skips_resolution(self) -> None:
        """Without entity_keys, the resolver is never called (zero overhead)."""
        client = _make_client()

        session = AsyncMock()
        mock_result = MagicMock()
        mock_result.scalars.return_value.all.return_value = []
        session.execute.return_value = mock_result
        client._session_factory = _mock_session_factory(session)

        with patch(
            "arandu.read.retrieval._resolve_entity_keys",
            new_callable=AsyncMock,
        ) as mock_resolve:
            await client.get_all(AGENT_ID)

            mock_resolve.assert_not_called()
