"""Tests for query expansion — entity priming + deterministic entity resolution."""

from __future__ import annotations

import json
from unittest.mock import AsyncMock, MagicMock

from conftest import run_async

from arandu.read.query_expansion import (
    ExpandedQuery,
    _extract_query_words,
    resolve_query_entities,
)
from arandu.read.retrieval_agent import _normalize_entities, _parse_plan

# ---------------------------------------------------------------------------
# Word extraction tests
# ---------------------------------------------------------------------------


def test_extract_query_words_basic():
    """Basic word extraction."""
    words = _extract_query_words("como está o Kevin?")
    assert "kevin" in words
    assert "como" in words


def test_extract_query_words_short():
    """Words shorter than 2 chars should be excluded."""
    words = _extract_query_words("o a é")
    assert len(words) == 0


def test_extract_query_words_accents():
    """Words with accents should be extracted."""
    words = _extract_query_words("são Paulo é bonito")
    assert "são" in words or "paulo" in words


# ---------------------------------------------------------------------------
# ExpandedQuery dataclass
# ---------------------------------------------------------------------------


def test_expanded_query_defaults():
    """ExpandedQuery should have sane defaults."""
    eq = ExpandedQuery(primed_entities=[], temporal_range=None, expanded_terms=[])
    assert eq.primed_entities == []
    assert eq.temporal_range is None
    assert eq.expanded_terms == []


# ---------------------------------------------------------------------------
# Deterministic entity resolution: resolve_query_entities()
# ---------------------------------------------------------------------------


def _mock_session_for_resolution(
    alias_results: list[tuple[str, str]] | None = None,
    display_results: list[tuple[str,]] | None = None,
    slug_results: list[tuple[str,]] | None = None,
) -> AsyncMock:
    """Create a mock session that returns different results per query.

    The function is called 3 times (alias, display_name, slug queries).
    We detect which query is running by inspecting the call count.
    """
    call_count = 0
    all_results = [alias_results or [], display_results or [], slug_results or []]

    async def mock_execute(stmt, *args, **kwargs):
        nonlocal call_count
        idx = min(call_count, len(all_results) - 1)
        call_count += 1
        mock_result = MagicMock()
        mock_result.all.return_value = all_results[idx]
        return mock_result

    session = AsyncMock()
    session.execute = mock_execute
    return session


class TestResolveQueryEntities:
    def test_resolve_by_alias(self) -> None:
        """Alias match resolves canonical entity_key."""
        session = _mock_session_for_resolution(
            alias_results=[("carlos", "person:carlos")],
        )
        result = run_async(resolve_query_entities(session, "user_1", "Onde o Carlos mora?"))
        assert "person:carlos" in result

    def test_resolve_by_display_name(self) -> None:
        """Display name match resolves canonical_key."""
        session = _mock_session_for_resolution(
            alias_results=[],
            display_results=[("person:carlos",)],
        )
        result = run_async(resolve_query_entities(session, "user_1", "Onde o Carlos mora?"))
        assert "person:carlos" in result

    def test_resolve_by_slug(self) -> None:
        """Slug match resolves entity_key from MemoryFact entity_keys."""
        session = _mock_session_for_resolution(
            alias_results=[],
            display_results=[],
            slug_results=[("person:carlos",)],
        )
        result = run_async(resolve_query_entities(session, "user_1", "Onde o Carlos mora?"))
        assert "person:carlos" in result

    def test_resolve_union_dedup(self) -> None:
        """Multiple sources resolving same entity should not duplicate."""
        session = _mock_session_for_resolution(
            alias_results=[("carlos", "person:carlos")],
            display_results=[("person:carlos",)],
            slug_results=[("person:carlos",)],
        )
        result = run_async(resolve_query_entities(session, "user_1", "Onde o Carlos mora?"))
        assert result.count("person:carlos") == 1

    def test_resolve_empty_query(self) -> None:
        """Empty query returns no entities."""
        session = AsyncMock()
        result = run_async(resolve_query_entities(session, "user_1", ""))
        assert result == []

    def test_resolve_failsafe_on_error(self) -> None:
        """DB error returns empty list (fail-safe)."""
        session = AsyncMock()
        session.execute = AsyncMock(side_effect=RuntimeError("DB down"))
        result = run_async(resolve_query_entities(session, "user_1", "Onde o Carlos mora?"))
        assert result == []


# ---------------------------------------------------------------------------
# Entity normalization in _parse_plan
# ---------------------------------------------------------------------------


class TestNormalizeEntities:
    def test_valid_entity_unchanged(self) -> None:
        """Already valid entity_key passes through."""
        schema = [("person:carlos", "location"), ("person:pedro", "profession")]
        result = _normalize_entities(["person:carlos"], schema)
        assert result == ["person:carlos"]

    def test_malformed_entity_normalized(self) -> None:
        """entity:Carlos → person:carlos via slug match."""
        schema = [("person:carlos", "location"), ("organization:techcorp", "industry")]
        result = _normalize_entities(["entity:Carlos"], schema)
        assert "person:carlos" in result

    def test_bare_name_normalized(self) -> None:
        """Bare name 'carlos' → person:carlos via slug match."""
        schema = [("person:carlos", "location")]
        result = _normalize_entities(["carlos"], schema)
        assert "person:carlos" in result

    def test_unknown_entity_kept(self) -> None:
        """Entity not in schema is kept as-is."""
        schema = [("person:carlos", "location")]
        result = _normalize_entities(["person:unknown"], schema)
        assert "person:unknown" in result

    def test_case_insensitive(self) -> None:
        """Normalization is case-insensitive."""
        schema = [("person:carlos", "location")]
        result = _normalize_entities(["Person:Carlos"], schema)
        assert "person:carlos" in result

    def test_dedup(self) -> None:
        """Duplicate entities after normalization are deduped."""
        schema = [("person:carlos", "location")]
        result = _normalize_entities(["person:carlos", "entity:Carlos"], schema)
        assert result.count("person:carlos") == 1


class TestParsePlanNormalization:
    def test_parse_plan_normalizes_entities(self) -> None:
        """_parse_plan normalizes entities when schema_pairs provided."""
        schema = [("person:carlos", "location"), ("person:pedro", "profession")]
        raw = json.dumps(
            {
                "strategy": "similarity",
                "similarity_query": "Carlos location",
                "entities": ["entity:Carlos"],
                "max_facts": 20,
                "reason": "specific search",
            }
        )
        plan = _parse_plan(raw, "Onde o Carlos mora?", schema_pairs=schema)
        assert plan is not None
        assert "person:carlos" in plan.entities

    def test_parse_plan_without_schema_no_normalization(self) -> None:
        """Without schema_pairs, entities pass through unchanged."""
        raw = json.dumps(
            {
                "strategy": "similarity",
                "similarity_query": "Carlos location",
                "entities": ["entity:Carlos"],
                "max_facts": 20,
                "reason": "specific search",
            }
        )
        plan = _parse_plan(raw, "Onde o Carlos mora?")
        assert plan is not None
        assert "entity:Carlos" in plan.entities
