"""
tests/test_entity_suggest_merges_v2.py — Tests for improved suggest_merges() (Issue #80).

Covers:
  - Embedding pre-filter happy path (pair count reduced)
  - Embedding pre-filter graceful fallback on LLM failure
  - Cross-session co-reference detection
  - Confidence threshold filtering
  - suggest_merges() with < 2 entities returns []
  - CortexPlugin.suggest_merges() delegation
  - CortexPlugin.merge_entities() delegation
  - CortexPlugin.merge_entities() raises ValueError for same-id merge
  - CortexPlugin.suggest_merges() returns [] when resolver unavailable
  - CortexPlugin.merge_entities() returns None when resolver unavailable
"""
from __future__ import annotations

import json
import pytest
import pytest_asyncio
import numpy as np
from unittest.mock import AsyncMock, MagicMock, patch, PropertyMock

from cortex.types import (
    Entity,
    EntityAlias,
    LLMProvider,
    MergeSuggestion,
    StorageAdapter,
)
from cortex.core.entity_resolution import EntityResolver


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _entity(id: str, canonical_name: str, owner_id: str = "default") -> Entity:
    return Entity(
        id=id,
        owner_id=owner_id,
        canonical_name=canonical_name,
        entity_type="person",
    )


def _alias(entity_id: str, alias: str, id: str = None) -> EntityAlias:
    return EntityAlias(
        id=id or f"alias-{entity_id}",
        entity_id=entity_id,
        alias=alias,
        alias_type="name",
        confidence=1.0,
    )


def _make_storage(aliases=None, entity_map=None, merged_entity=None):
    storage = AsyncMock(spec=StorageAdapter)
    storage.get_all_entity_aliases.return_value = aliases or []
    entity_map = entity_map or {}

    async def get_entity(entity_id):
        return entity_map.get(entity_id)

    storage.get_entity_by_id.side_effect = get_entity
    storage.merge_entities.return_value = merged_entity or _entity("ent-keep", "Andrew")
    return storage


def _make_llm(embeddings=None, fail_embed=False):
    llm = AsyncMock(spec=LLMProvider)
    if fail_embed:
        llm.embed.side_effect = RuntimeError("embed API unavailable")
    else:
        llm.embed.return_value = embeddings or []
    return llm


def _unit_vec(v):
    """Return unit vector for a list."""
    arr = np.array(v, dtype=np.float32)
    norm = np.linalg.norm(arr)
    return (arr / norm if norm > 0 else arr).tolist()


# ---------------------------------------------------------------------------
# 1. Embedding pre-filter — happy path
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_embedding_prefilter_reduces_pairs():
    """Entity pairs with low cosine similarity must be skipped by the pre-filter."""
    # Three entities: A ≈ B (similar embeddings), C far from both
    ent_a = _entity("ent-a", "Andrew")
    ent_b = _entity("ent-b", "Androw")  # typo of Andrew
    ent_c = _entity("ent-c", "Zephyr")

    aliases = [
        _alias("ent-a", "Andrew"),
        _alias("ent-b", "Androw"),
        _alias("ent-c", "Zephyr"),
    ]
    entity_map = {"ent-a": ent_a, "ent-b": ent_b, "ent-c": ent_c}

    storage = _make_storage(aliases=aliases, entity_map=entity_map)

    # Construct embeddings: A and B very similar, C orthogonal
    vec_a = _unit_vec([1.0, 0.1, 0.0])
    vec_b = _unit_vec([0.95, 0.1, 0.0])   # very close to A
    vec_c = _unit_vec([0.0, 0.0, 1.0])    # orthogonal to A and B
    llm = _make_llm(embeddings=[vec_a, vec_b, vec_c])

    resolver = EntityResolver(storage=storage, llm=llm)

    # With prefilter threshold=0.80, only A-B pair should pass (C filtered out)
    suggestions = await resolver.suggest_merges(
        owner_id="default",
        similarity_threshold=50.0,
        confidence_threshold=0.0,
        use_embedding_prefilter=True,
        embedding_prefilter_threshold=0.80,
    )

    # get_all_entity_aliases was called once
    storage.get_all_entity_aliases.assert_called_once_with(owner_id="default")
    # embed was called with canonical names
    llm.embed.assert_called_once()

    # We should have a suggestion for A-B
    assert len(suggestions) >= 1
    pair_ids = {(s.keep_id, s.merge_id) for s in suggestions}
    assert ("ent-a", "ent-b") in pair_ids or ("ent-b", "ent-a") in pair_ids

    # C should NOT appear in any suggestion
    for s in suggestions:
        assert "ent-c" not in (s.keep_id, s.merge_id), (
            "Zephyr entity should have been filtered by embedding pre-filter"
        )


# ---------------------------------------------------------------------------
# 2. Embedding pre-filter — fallback on failure
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_embedding_prefilter_fallback_on_llm_failure():
    """When embed() raises, suggest_merges() gracefully falls back to full O(n²)."""
    ent_a = _entity("ent-a", "Andrew")
    ent_b = _entity("ent-b", "Androw")
    aliases = [_alias("ent-a", "Andrew"), _alias("ent-b", "Androw")]
    entity_map = {"ent-a": ent_a, "ent-b": ent_b}

    storage = _make_storage(aliases=aliases, entity_map=entity_map)
    llm = _make_llm(fail_embed=True)

    resolver = EntityResolver(storage=storage, llm=llm)

    # Should not raise — should fall back to O(n²) and still find the pair
    suggestions = await resolver.suggest_merges(
        similarity_threshold=70.0,
        confidence_threshold=0.0,
        use_embedding_prefilter=True,
    )

    assert len(suggestions) >= 1
    pair_ids = {(s.keep_id, s.merge_id) for s in suggestions}
    assert ("ent-a", "ent-b") in pair_ids or ("ent-b", "ent-a") in pair_ids


# ---------------------------------------------------------------------------
# 3. Cross-session co-reference detection
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_cross_session_coref_finds_same_entity_across_owners():
    """Entities from different owner IDs (sessions) should be suggested for merge."""
    ent_a = _entity("ent-a", "Andrew", owner_id="session-1")
    ent_b = _entity("ent-b", "Andrew", owner_id="session-2")

    aliases_s1 = [_alias("ent-a", "Andrew")]
    aliases_s2 = [_alias("ent-b", "Andrew")]
    entity_map = {"ent-a": ent_a, "ent-b": ent_b}

    storage = AsyncMock(spec=StorageAdapter)

    async def get_all_aliases(owner_id="default"):
        if owner_id == "session-1":
            return aliases_s1
        if owner_id == "session-2":
            return aliases_s2
        return []

    storage.get_all_entity_aliases.side_effect = get_all_aliases

    async def get_entity(entity_id):
        return entity_map.get(entity_id)

    storage.get_entity_by_id.side_effect = get_entity

    # No embedding — use pure fuzzy (identical names → score 100)
    llm = _make_llm(embeddings=[_unit_vec([1, 0]), _unit_vec([1, 0])])

    resolver = EntityResolver(storage=storage, llm=llm)

    suggestions = await resolver.suggest_merges(
        owner_id="session-1",
        similarity_threshold=80.0,
        confidence_threshold=0.0,
        cross_session_owner_ids=["session-2"],
        use_embedding_prefilter=False,  # keep it simple for this test
    )

    assert len(suggestions) == 1
    s = suggestions[0]
    # Should flag as cross-session in the reason
    assert "cross-session" in s.reason.lower() or "session" in s.reason.lower()
    assert {s.keep_id, s.merge_id} == {"ent-a", "ent-b"}
    assert s.similarity_score == 1.0


# ---------------------------------------------------------------------------
# 4. Confidence threshold filtering
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_confidence_threshold_filters_low_confidence():
    """Suggestions with normalised score below confidence_threshold are excluded."""
    ent_a = _entity("ent-a", "Andrew")
    ent_b = _entity("ent-b", "Androw")
    aliases = [_alias("ent-a", "Andrew"), _alias("ent-b", "Androw")]
    entity_map = {"ent-a": ent_a, "ent-b": ent_b}

    storage = _make_storage(aliases=aliases, entity_map=entity_map)
    # Disable embedding prefilter for simplicity
    llm = _make_llm(fail_embed=True)

    resolver = EntityResolver(storage=storage, llm=llm)

    # Andrew vs Androw → token_sort_ratio ≈ 83 → normalised ≈ 0.83
    # With confidence_threshold=0.90, it should be excluded
    high_threshold_results = await resolver.suggest_merges(
        similarity_threshold=70.0,
        confidence_threshold=0.90,
        use_embedding_prefilter=False,
    )
    assert high_threshold_results == []

    # With confidence_threshold=0.50, it should be included
    low_threshold_results = await resolver.suggest_merges(
        similarity_threshold=70.0,
        confidence_threshold=0.50,
        use_embedding_prefilter=False,
    )
    assert len(low_threshold_results) >= 1


# ---------------------------------------------------------------------------
# 5. Less than 2 entities returns empty list
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_suggest_merges_single_entity_returns_empty():
    """With only one entity, there are no pairs to compare."""
    storage = _make_storage(
        aliases=[_alias("ent-a", "Andrew")],
        entity_map={"ent-a": _entity("ent-a", "Andrew")},
    )
    llm = _make_llm()
    resolver = EntityResolver(storage=storage, llm=llm)

    suggestions = await resolver.suggest_merges()
    assert suggestions == []


# ---------------------------------------------------------------------------
# 6. CortexPlugin.suggest_merges() — delegation
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_plugin_suggest_merges_delegates_to_resolver():
    """CortexPlugin.suggest_merges() must delegate to entity_resolver."""
    from cortex.api.plugin import CortexPlugin

    plugin = MagicMock(spec=CortexPlugin)
    plugin.config = MagicMock()
    plugin.config.owner_id = "test-owner"

    expected = [
        MergeSuggestion(
            keep_id="ent-a",
            merge_id="ent-b",
            keep_name="Andrew",
            merge_name="Androw",
            similarity_score=0.83,
            reason="Alias fuzzy similarity 83.0/100",
        )
    ]

    mock_resolver = AsyncMock()
    mock_resolver.suggest_merges.return_value = expected
    plugin.entity_resolver = mock_resolver

    # Call the real plugin method (unbound, passing plugin as self)
    result = await CortexPlugin.suggest_merges(
        plugin,
        similarity_threshold=75.0,
        confidence_threshold=0.5,
        use_embedding_prefilter=True,
    )

    mock_resolver.suggest_merges.assert_called_once_with(
        owner_id="test-owner",
        similarity_threshold=75.0,
        confidence_threshold=0.5,
        cross_session_owner_ids=None,
        use_embedding_prefilter=True,
        embedding_prefilter_threshold=0.60,
    )
    assert result == expected


# ---------------------------------------------------------------------------
# 7. CortexPlugin.suggest_merges() — resolver unavailable
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_plugin_suggest_merges_returns_empty_when_no_resolver():
    """plugin.suggest_merges() returns [] gracefully when EntityResolver is absent."""
    from cortex.api.plugin import CortexPlugin

    plugin = MagicMock(spec=CortexPlugin)
    plugin.entity_resolver = None
    plugin.config = MagicMock()
    plugin.config.owner_id = "default"

    result = await CortexPlugin.suggest_merges(plugin)
    assert result == []


# ---------------------------------------------------------------------------
# 8. CortexPlugin.merge_entities() — delegation
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_plugin_merge_entities_delegates_to_resolver():
    """CortexPlugin.merge_entities() must delegate to entity_resolver.merge_entities()."""
    from cortex.api.plugin import CortexPlugin

    plugin = MagicMock(spec=CortexPlugin)
    kept = _entity("ent-keep", "Andrew")
    mock_resolver = AsyncMock()
    mock_resolver.merge_entities.return_value = kept
    plugin.entity_resolver = mock_resolver

    result = await CortexPlugin.merge_entities(plugin, keep_id="ent-keep", merge_id="ent-dup")

    mock_resolver.merge_entities.assert_called_once_with(
        keep_id="ent-keep",
        merge_id="ent-dup",
    )
    assert result.id == "ent-keep"


# ---------------------------------------------------------------------------
# 9. CortexPlugin.merge_entities() — same-id raises ValueError
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_plugin_merge_entities_propagates_value_error():
    """CortexPlugin.merge_entities() must re-raise ValueError for same-id merge."""
    from cortex.api.plugin import CortexPlugin

    plugin = MagicMock(spec=CortexPlugin)
    mock_resolver = AsyncMock()
    mock_resolver.merge_entities.side_effect = ValueError("keep_id and merge_id must be different")
    plugin.entity_resolver = mock_resolver

    with pytest.raises(ValueError, match="different"):
        await CortexPlugin.merge_entities(plugin, keep_id="ent-x", merge_id="ent-x")


# ---------------------------------------------------------------------------
# 10. CortexPlugin.merge_entities() — resolver unavailable
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_plugin_merge_entities_returns_none_when_no_resolver():
    """plugin.merge_entities() returns None gracefully when EntityResolver is absent."""
    from cortex.api.plugin import CortexPlugin

    plugin = MagicMock(spec=CortexPlugin)
    plugin.entity_resolver = None

    result = await CortexPlugin.merge_entities(plugin, keep_id="ent-a", merge_id="ent-b")
    assert result is None


# ---------------------------------------------------------------------------
# 11. Cross-session: same owner excluded from double-counting
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_cross_session_skips_primary_owner_id():
    """When primary owner_id is repeated in cross_session_owner_ids, aliases are not double-loaded."""
    ent_a = _entity("ent-a", "Andrew")
    aliases = [_alias("ent-a", "Andrew")]
    entity_map = {"ent-a": ent_a}

    storage = _make_storage(aliases=aliases, entity_map=entity_map)
    llm = _make_llm()
    resolver = EntityResolver(storage=storage, llm=llm)

    # Passing the same owner_id in cross_session_owner_ids should not duplicate
    suggestions = await resolver.suggest_merges(
        owner_id="default",
        cross_session_owner_ids=["default"],  # same as primary — should be skipped
        use_embedding_prefilter=False,
    )
    # Only one entity → no pairs → no suggestions
    assert suggestions == []
    # get_all_entity_aliases called only once (the duplicate "default" skipped)
    storage.get_all_entity_aliases.assert_called_once_with(owner_id="default")
