"""
tests/test_entity_resolution.py — Tests for EntityResolver (Module 3)

All storage and LLM calls are mocked. Tests cover:
  - Exact alias match
  - Fuzzy alias match (above threshold)
  - Ambiguous fuzzy match → LLM disambiguation
  - No match → new entity creation
  - merge_entities()
  - add_alias()
  - suggest_merges()
  - resolve_batch()
  - Edge cases: empty alias store, empty mention, duplicate keep/merge ids
"""
from __future__ import annotations

import json
import pytest
import pytest_asyncio
from unittest.mock import AsyncMock, MagicMock, patch

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


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


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


def _alias(
    id: str = "alias-001",
    entity_id: str = "ent-001",
    alias: str = "Andrew",
    alias_type: str = "name",
    confidence: float = 1.0,
) -> EntityAlias:
    return EntityAlias(
        id=id,
        entity_id=entity_id,
        alias=alias,
        alias_type=alias_type,
        confidence=confidence,
    )


def _make_storage(
    find_result: Entity | None = None,
    aliases: list[EntityAlias] | None = None,
    entity_by_id: Entity | None = None,
    created_entity: Entity | None = None,
    merged_entity: Entity | None = None,
    added_alias: EntityAlias | None = None,
    claims: list[SemanticClaim] | None = None,
) -> StorageAdapter:
    storage = AsyncMock(spec=StorageAdapter)
    storage.find_entity_by_alias.return_value = find_result
    storage.get_all_entity_aliases.return_value = aliases or []
    storage.get_entity_by_id.return_value = entity_by_id
    storage.create_entity.return_value = created_entity or _entity()
    storage.merge_entities.return_value = merged_entity or _entity()
    storage.add_entity_alias.return_value = added_alias or _alias()
    storage.get_claims_by_entity.return_value = claims or []
    return storage


def _make_llm(response_json: str | None = None) -> LLMProvider:
    llm = AsyncMock(spec=LLMProvider)
    llm.complete.return_value = response_json or json.dumps({
        "best_match_index": 0,
        "confidence": 0.9,
        "reasoning": "Clear match",
    })
    return llm


# ---------------------------------------------------------------------------
# Fixtures
# ---------------------------------------------------------------------------


@pytest.fixture
def resolver_no_entities():
    """Resolver with an empty storage (no entities)."""
    storage = _make_storage(find_result=None, aliases=[], entity_by_id=None)
    llm = _make_llm()
    return EntityResolver(storage=storage, llm=llm)


@pytest.fixture
def andrew_entity():
    return _entity(id="ent-andrew", canonical_name="Andrew", entity_type="person")


@pytest.fixture
def andrew_alias(andrew_entity):
    return _alias(id="alias-andrew", entity_id=andrew_entity.id, alias="Andrew")


# ---------------------------------------------------------------------------
# 1. Exact alias match
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_exact_match_returns_entity(andrew_entity, andrew_alias):
    """Exact case-insensitive alias match skips LLM entirely."""
    storage = _make_storage(find_result=andrew_entity)
    llm = _make_llm()
    resolver = EntityResolver(storage=storage, llm=llm)

    result = await resolver.resolve("Andrew", owner_id="default")

    assert result.entity_id == andrew_entity.id
    assert result.canonical_name == "Andrew"
    assert result.confidence == 1.0
    assert result.is_new is False
    assert result.matched_alias == "Andrew"

    # LLM must NOT be called for an exact match
    llm.complete.assert_not_called()
    storage.find_entity_by_alias.assert_called_once_with("Andrew", owner_id="default")


@pytest.mark.asyncio
async def test_exact_match_case_insensitive_delegated_to_storage(andrew_entity):
    """Storage handles case-insensitive lookup; resolver trusts its result."""
    storage = _make_storage(find_result=andrew_entity)
    llm = _make_llm()
    resolver = EntityResolver(storage=storage, llm=llm)

    result = await resolver.resolve("andrew")
    # Storage returned the entity — resolver should trust it
    assert result.entity_id == andrew_entity.id
    assert result.is_new is False


# ---------------------------------------------------------------------------
# 2. Fuzzy alias match (above threshold, unambiguous)
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_fuzzy_match_above_threshold(andrew_entity, andrew_alias):
    """'Andew' is a clear typo for 'Andrew' — score ≥ 85, unambiguous → match."""
    storage = _make_storage(
        find_result=None,          # no exact match
        aliases=[andrew_alias],
        entity_by_id=andrew_entity,
    )
    llm = _make_llm()
    resolver = EntityResolver(storage=storage, llm=llm, fuzzy_threshold=85.0)

    result = await resolver.resolve("Andew")  # one character off

    assert result.entity_id == andrew_entity.id
    assert result.is_new is False
    assert result.confidence >= 0.60


@pytest.mark.asyncio
async def test_fuzzy_match_below_threshold_creates_new_entity():
    """A very dissimilar mention should NOT match and create a new entity instead."""
    some_entity = _entity(id="ent-xyz", canonical_name="Zephyr")
    alias = _alias(entity_id="ent-xyz", alias="Zephyr")
    new_entity = _entity(id="ent-new", canonical_name="Completely Different Thing")

    storage = _make_storage(
        find_result=None,
        aliases=[alias],
        entity_by_id=some_entity,
        created_entity=new_entity,
    )
    llm = AsyncMock(spec=LLMProvider)
    # LLM should say no match (null best_match_index)
    llm.complete.return_value = json.dumps({
        "best_match_index": None,
        "confidence": 0.1,
        "reasoning": "No good match",
    })
    resolver = EntityResolver(storage=storage, llm=llm, fuzzy_threshold=85.0)

    result = await resolver.resolve("Completely Different Thing")
    assert result.is_new is True


# ---------------------------------------------------------------------------
# 3. New entity creation (no match at all)
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_no_aliases_creates_new_entity():
    """When the store is empty, resolve() always creates a new entity."""
    new_entity = _entity(id="ent-new", canonical_name="Eva")
    storage = _make_storage(find_result=None, aliases=[], created_entity=new_entity)
    llm = _make_llm()
    resolver = EntityResolver(storage=storage, llm=llm)

    result = await resolver.resolve("Eva", owner_id="default")

    assert result.is_new is True
    assert result.entity_id == "ent-new"
    assert result.canonical_name == "Eva"
    assert result.confidence == 1.0
    storage.create_entity.assert_called_once_with(
        canonical_name="Eva",
        entity_type="unknown",
        owner_id="default",
    )


@pytest.mark.asyncio
async def test_create_entity_adds_initial_alias():
    """create_entity() must add the canonical name as the entity's first alias."""
    entity = _entity(id="ent-lucca", canonical_name="Lucca")
    storage = _make_storage(created_entity=entity)
    llm = _make_llm()
    resolver = EntityResolver(storage=storage, llm=llm)

    await resolver.create_entity("Lucca", entity_type="person")

    storage.add_entity_alias.assert_called_once_with(
        entity_id="ent-lucca",
        alias="Lucca",
        alias_type="name",
        source="operator",
        confidence=1.0,
    )


@pytest.mark.asyncio
async def test_empty_mention_raises_value_error():
    """resolve() must reject empty strings immediately."""
    resolver = EntityResolver(storage=AsyncMock(), llm=AsyncMock())
    with pytest.raises(ValueError, match="non-empty"):
        await resolver.resolve("  ")


# ---------------------------------------------------------------------------
# 4. LLM disambiguation (ambiguous fuzzy matches)
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_llm_disambiguation_picks_correct_candidate():
    """When two aliases both score ≥60 and are within 5 pts of each other,
    the LLM is asked to disambiguate and its choice is respected."""
    entity_a = _entity(id="ent-a", canonical_name="Andrew Smith")
    entity_b = _entity(id="ent-b", canonical_name="Andrew Jones")
    alias_a = _alias(id="al-a", entity_id="ent-a", alias="Andrew Smith")
    alias_b = _alias(id="al-b", entity_id="ent-b", alias="Andrew Jones")

    async def get_entity(entity_id: str):
        return entity_a if entity_id == "ent-a" else entity_b

    storage = _make_storage(
        find_result=None,
        aliases=[alias_a, alias_b],
    )
    storage.get_entity_by_id.side_effect = get_entity

    llm = AsyncMock(spec=LLMProvider)
    llm.complete.return_value = json.dumps({
        "best_match_index": 0,
        "confidence": 0.88,
        "reasoning": "Context implies Smith not Jones",
    })
    resolver = EntityResolver(storage=storage, llm=llm, fuzzy_threshold=85.0)

    result = await resolver.resolve("Andrew", context="Andrew Smith is in the meeting")

    # LLM was consulted
    llm.complete.assert_called_once()
    assert result.entity_id in ("ent-a", "ent-b")


@pytest.mark.asyncio
async def test_llm_returns_no_match_triggers_new_entity():
    """When the LLM says no candidate matches, create a new entity."""
    entity_a = _entity(id="ent-a", canonical_name="Andrew")
    alias_a = _alias(id="al-a", entity_id="ent-a", alias="Andrew")
    new_entity = _entity(id="ent-new", canonical_name="Andrius")

    storage = _make_storage(
        find_result=None,
        aliases=[alias_a],
        entity_by_id=entity_a,
        created_entity=new_entity,
    )
    llm = AsyncMock(spec=LLMProvider)
    llm.complete.return_value = json.dumps({
        "best_match_index": None,
        "confidence": 0.3,
        "reasoning": "Different person entirely",
    })
    resolver = EntityResolver(storage=storage, llm=llm, fuzzy_threshold=85.0)

    result = await resolver.resolve("Andrius")
    assert result.is_new is True
    assert result.entity_id == "ent-new"


@pytest.mark.asyncio
async def test_llm_failure_falls_back_to_new_entity(andrew_entity, andrew_alias):
    """If the LLM call throws an exception, fall back to new entity creation.

    Use fuzzy_threshold=95 so 'Andreew' (scores ~92% against 'Andrew')
    falls in the ambiguous zone — triggering an LLM call that then fails.
    """
    new_entity = _entity(id="ent-new", canonical_name="Andreew")

    storage = _make_storage(
        find_result=None,
        aliases=[andrew_alias],
        entity_by_id=andrew_entity,
        created_entity=new_entity,
    )
    llm = AsyncMock(spec=LLMProvider)
    llm.complete.side_effect = RuntimeError("LLM unavailable")
    # threshold=95 → score 92 < 95 → goes to LLM → LLM throws → new entity
    resolver = EntityResolver(storage=storage, llm=llm, fuzzy_threshold=95.0)

    result = await resolver.resolve("Andreew")
    assert result.is_new is True


# ---------------------------------------------------------------------------
# 5. merge_entities()
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_merge_delegates_to_storage():
    """merge_entities() must delegate to storage.merge_entities()."""
    merged = _entity(id="ent-keep", canonical_name="Andrew")
    storage = _make_storage(merged_entity=merged)
    llm = _make_llm()
    resolver = EntityResolver(storage=storage, llm=llm)

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

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


@pytest.mark.asyncio
async def test_merge_same_id_raises():
    """Merging an entity with itself must raise ValueError."""
    resolver = EntityResolver(storage=AsyncMock(), llm=AsyncMock())
    with pytest.raises(ValueError, match="different"):
        await resolver.merge_entities("ent-001", "ent-001")


# ---------------------------------------------------------------------------
# 6. add_alias()
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_add_alias_delegates_to_storage():
    alias = _alias(entity_id="ent-andrew", alias="Andy")
    storage = _make_storage(added_alias=alias)
    llm = _make_llm()
    resolver = EntityResolver(storage=storage, llm=llm)

    result = await resolver.add_alias(entity_id="ent-andrew", alias="Andy")

    storage.add_entity_alias.assert_called_once_with(
        entity_id="ent-andrew",
        alias="Andy",
        alias_type="name",
        source="extracted",
        confidence=1.0,
    )
    assert result.alias == "Andy"


# ---------------------------------------------------------------------------
# 7. suggest_merges()
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_suggest_merges_finds_similar_entities():
    """Two entities with near-identical aliases should be suggested for merge.

    'Andrew' vs 'Androw' → token_sort_ratio ≈ 83 which is above threshold=75.
    """
    entity_a = _entity(id="ent-a", canonical_name="Andrew")
    entity_b = _entity(id="ent-b", canonical_name="Androw")
    aliases = [
        _alias(id="al-a", entity_id="ent-a", alias="Andrew"),
        _alias(id="al-b", entity_id="ent-b", alias="Androw"),
    ]

    async def get_entity(entity_id: str):
        return entity_a if entity_id == "ent-a" else entity_b

    storage = _make_storage(aliases=aliases)
    storage.get_entity_by_id.side_effect = get_entity
    llm = _make_llm()
    resolver = EntityResolver(storage=storage, llm=llm)

    suggestions = await resolver.suggest_merges(owner_id="default", similarity_threshold=75.0)

    assert len(suggestions) >= 1
    assert any(
        (s.keep_id == "ent-a" and s.merge_id == "ent-b")
        or (s.keep_id == "ent-b" and s.merge_id == "ent-a")
        for s in suggestions
    )


@pytest.mark.asyncio
async def test_suggest_merges_empty_store():
    """No entities → no suggestions."""
    storage = _make_storage(aliases=[])
    resolver = EntityResolver(storage=storage, llm=_make_llm())

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


@pytest.mark.asyncio
async def test_suggest_merges_no_similar_entities():
    """Very different entities should not be suggested for merge."""
    aliases = [
        _alias(id="al-a", entity_id="ent-a", alias="Zephyr"),
        _alias(id="al-b", entity_id="ent-b", alias="Tokyo"),
    ]

    async def get_entity(entity_id: str):
        return _entity(id=entity_id, canonical_name=entity_id)

    storage = _make_storage(aliases=aliases)
    storage.get_entity_by_id.side_effect = get_entity
    resolver = EntityResolver(storage=storage, llm=_make_llm())

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


# ---------------------------------------------------------------------------
# 8. resolve_batch()
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_resolve_batch_resolves_all_mentions():
    """resolve_batch() returns one ResolvedEntity per input CandidateEntity."""
    entity_a = _entity(id="ent-a", canonical_name="Andrew")
    entity_b = _entity(id="ent-b", canonical_name="Eva")

    call_count = 0
    entities = [entity_a, entity_b]

    async def find_by_alias(alias, owner_id="default"):
        nonlocal call_count
        result = entities[call_count] if call_count < len(entities) else None
        call_count += 1
        return result

    storage = AsyncMock(spec=StorageAdapter)
    storage.find_entity_by_alias.side_effect = find_by_alias
    storage.get_all_entity_aliases.return_value = []
    storage.create_entity.return_value = _entity(id="ent-new", canonical_name="New")
    storage.add_entity_alias.return_value = _alias()

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

    mentions = [
        CandidateEntity(name="Andrew", entity_type="person"),
        CandidateEntity(name="Eva", entity_type="person"),
    ]
    results = await resolver.resolve_batch(mentions, owner_id="default")

    assert len(results) == 2
    assert results[0].entity_id == "ent-a"
    assert results[1].entity_id == "ent-b"


# ---------------------------------------------------------------------------
# 9. Entity type upgrade (unknown → typed)
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_exact_match_upgrades_unknown_type():
    """When an existing entity has type 'unknown' and resolve() is called with
    a real entity_type, the entity type should be upgraded in storage."""
    entity = _entity(id="ent-u", canonical_name="Cortex", entity_type="unknown")
    storage = _make_storage(find_result=entity)
    storage.update_entity = AsyncMock()
    llm = _make_llm()
    resolver = EntityResolver(storage=storage, llm=llm)

    result = await resolver.resolve("Cortex", entity_type="project", owner_id="default")

    assert result.entity_id == "ent-u"
    assert result.entity_type == "project"
    assert result.is_new is False
    storage.update_entity.assert_called_once_with(
        "ent-u", entity_type="project", owner_id="default",
    )


@pytest.mark.asyncio
async def test_exact_match_does_not_downgrade_typed_entity():
    """If entity already has a real type, passing 'unknown' should NOT overwrite it."""
    entity = _entity(id="ent-t", canonical_name="Python", entity_type="technology")
    storage = _make_storage(find_result=entity)
    storage.update_entity = AsyncMock()
    llm = _make_llm()
    resolver = EntityResolver(storage=storage, llm=llm)

    result = await resolver.resolve("Python", entity_type="unknown", owner_id="default")

    assert result.entity_type == "technology"
    storage.update_entity.assert_not_called()


@pytest.mark.asyncio
async def test_exact_match_no_upgrade_when_both_typed():
    """If entity already has a real type, don't overwrite even with a different real type."""
    entity = _entity(id="ent-t2", canonical_name="React", entity_type="technology")
    storage = _make_storage(find_result=entity)
    storage.update_entity = AsyncMock()
    llm = _make_llm()
    resolver = EntityResolver(storage=storage, llm=llm)

    result = await resolver.resolve("React", entity_type="project", owner_id="default")

    assert result.entity_type == "technology"  # original type preserved
    storage.update_entity.assert_not_called()


# ---------------------------------------------------------------------------
# 10. Configurable threshold
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_custom_fuzzy_threshold_respected():
    """Setting threshold=100 means only exact-score matches are accepted."""
    entity = _entity(id="ent-a", canonical_name="Andrew")
    alias = _alias(entity_id="ent-a", alias="Andrew")
    new_entity = _entity(id="ent-new", canonical_name="Androw")

    storage = _make_storage(
        find_result=None,
        aliases=[alias],
        entity_by_id=entity,
        created_entity=new_entity,
    )
    llm = AsyncMock(spec=LLMProvider)
    llm.complete.return_value = json.dumps({
        "best_match_index": None,
        "confidence": 0.4,
        "reasoning": "threshold 100 set",
    })
    # threshold=100 → no fuzzy match accepted → LLM says no → new entity
    resolver = EntityResolver(storage=storage, llm=llm, fuzzy_threshold=100.0)

    result = await resolver.resolve("Androw")
    assert result.is_new is True


# ---------------------------------------------------------------------------
# Entity type upgrade tests
# ---------------------------------------------------------------------------

@pytest.mark.asyncio
async def test_exact_match_upgrades_unknown_type():
    """When an existing entity has type 'unknown' and resolve() is called with
    a real entity_type, the entity type should be upgraded in storage."""
    entity = _entity(id="ent-u", canonical_name="Cortex", entity_type="unknown")
    storage = _make_storage(find_result=entity)
    storage.update_entity = AsyncMock()
    llm = _make_llm()
    resolver = EntityResolver(storage=storage, llm=llm)

    result = await resolver.resolve("Cortex", entity_type="project", owner_id="default")

    assert result.entity_id == "ent-u"
    assert result.entity_type == "project"
    assert result.is_new is False
    storage.update_entity.assert_called_once_with(
        "ent-u", entity_type="project", owner_id="default",
    )


@pytest.mark.asyncio
async def test_exact_match_does_not_downgrade_typed_entity():
    """If entity already has a real type, passing 'unknown' should NOT overwrite it."""
    entity = _entity(id="ent-t", canonical_name="Python", entity_type="technology")
    storage = _make_storage(find_result=entity)
    storage.update_entity = AsyncMock()
    llm = _make_llm()
    resolver = EntityResolver(storage=storage, llm=llm)

    result = await resolver.resolve("Python", entity_type="unknown", owner_id="default")

    assert result.entity_type == "technology"
    storage.update_entity.assert_not_called()


@pytest.mark.asyncio
async def test_exact_match_no_upgrade_when_both_typed():
    """If entity already has a real type, don't overwrite even with a different real type."""
    entity = _entity(id="ent-t2", canonical_name="React", entity_type="technology")
    storage = _make_storage(find_result=entity)
    storage.update_entity = AsyncMock()
    llm = _make_llm()
    resolver = EntityResolver(storage=storage, llm=llm)

    result = await resolver.resolve("React", entity_type="project", owner_id="default")

    assert result.entity_type == "technology"  # original type preserved
    storage.update_entity.assert_not_called()
