"""
tests/test_entity_resolution_audit.py — Tests for entity resolution audit fixes.

Covers:
  - B-001: N guard on dense N×N cosine matrix (OOM protection)
  - B-003: argpartition boundary guard when k >= array size
  - B-004: _load_aliases_paginated renamed to _load_all_aliases
  - B-017: argpartition deterministic ordering at ties
  - B-030: candidate_k=0 early return instead of O(n²) fallthrough
"""
from __future__ import annotations

import numpy as np
import pytest
from unittest.mock import AsyncMock, MagicMock

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


# ---------------------------------------------------------------------------
# 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 = "") -> EntityAlias:
    return EntityAlias(
        id=id or f"alias-{entity_id}-{alias}",
        entity_id=entity_id,
        alias=alias,
        alias_type="name",
        confidence=1.0,
    )


def _make_storage(aliases=None, entity_map=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
    # Ensure no pagination support so _load_all_aliases uses fallback
    if hasattr(storage, "get_entity_aliases_page"):
        del storage.get_entity_aliases_page
    return storage


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


# ---------------------------------------------------------------------------
# B-001: N guard on dense N×N cosine matrix → OOM at scale
# ---------------------------------------------------------------------------


class TestB001NGuard:
    """Verify that suggest_merges raises ValueError when N exceeds max."""

    @pytest.mark.asyncio
    async def test_exceeds_max_entities_raises(self):
        """N > max_entities_for_matrix should raise ValueError."""
        n = 15  # Use small N but set max to even smaller
        aliases = []
        entity_map = {}
        for i in range(n):
            eid = f"ent-{i}"
            name = f"Entity{i}"
            aliases.append(_alias(eid, name))
            entity_map[eid] = _entity(eid, name)

        storage = _make_storage(aliases=aliases, entity_map=entity_map)
        # Create embeddings that would work if not guarded
        embeddings = np.random.randn(n, 8).tolist()
        llm = _make_llm(embeddings=embeddings)

        resolver = EntityResolver(
            storage=storage,
            llm=llm,
            candidate_k=5,
            max_entities_for_matrix=10,  # Set max below N
        )

        with pytest.raises(ValueError, match="exceeds max_entities_for_matrix"):
            await resolver.suggest_merges(owner_id="default")

    @pytest.mark.asyncio
    async def test_within_max_entities_succeeds(self):
        """N <= max_entities_for_matrix should not raise."""
        n = 5
        aliases = []
        entity_map = {}
        for i in range(n):
            eid = f"ent-{i}"
            name = f"Entity{i}"
            aliases.append(_alias(eid, name))
            entity_map[eid] = _entity(eid, name)

        storage = _make_storage(aliases=aliases, entity_map=entity_map)
        embeddings = np.random.randn(n, 8).tolist()
        llm = _make_llm(embeddings=embeddings)

        resolver = EntityResolver(
            storage=storage,
            llm=llm,
            candidate_k=3,
            max_entities_for_matrix=100,  # Well above N
        )

        # Should not raise
        result = await resolver.suggest_merges(owner_id="default")
        assert isinstance(result, list)

    @pytest.mark.asyncio
    async def test_default_max_is_50000(self):
        """Default max_entities_for_matrix should be 50000."""
        storage = AsyncMock(spec=StorageAdapter)
        llm = AsyncMock(spec=LLMProvider)
        resolver = EntityResolver(storage=storage, llm=llm)
        assert resolver.max_entities_for_matrix == 50_000


# ---------------------------------------------------------------------------
# B-003: argpartition boundary fragile
# ---------------------------------------------------------------------------


class TestB003ArgpartitionBoundary:
    """Verify argpartition works when k >= array size."""

    @pytest.mark.asyncio
    async def test_candidate_k_exceeds_entity_count(self):
        """candidate_k > n-1 should not crash argpartition."""
        n = 3  # Only 3 entities, but candidate_k=100
        aliases = []
        entity_map = {}
        for i in range(n):
            eid = f"ent-{i}"
            name = f"Name{i}"
            aliases.append(_alias(eid, name))
            entity_map[eid] = _entity(eid, name)

        storage = _make_storage(aliases=aliases, entity_map=entity_map)
        # Create distinct embeddings
        embeddings = [[1, 0, 0], [0, 1, 0], [0, 0, 1]]
        llm = _make_llm(embeddings=embeddings)

        resolver = EntityResolver(
            storage=storage,
            llm=llm,
            candidate_k=100,  # Much larger than n-1=2
        )

        # Should not crash
        result = await resolver.suggest_merges(owner_id="default")
        assert isinstance(result, list)

    @pytest.mark.asyncio
    async def test_exactly_two_entities(self):
        """Edge case: exactly 2 entities, candidate_k=10."""
        aliases = [
            _alias("ent-a", "Alice"),
            _alias("ent-b", "Bob"),
        ]
        entity_map = {
            "ent-a": _entity("ent-a", "Alice"),
            "ent-b": _entity("ent-b", "Bob"),
        }

        storage = _make_storage(aliases=aliases, entity_map=entity_map)
        embeddings = [[1, 0], [0, 1]]
        llm = _make_llm(embeddings=embeddings)

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

        result = await resolver.suggest_merges(owner_id="default")
        assert isinstance(result, list)


# ---------------------------------------------------------------------------
# B-004: _load_aliases_paginated renamed to _load_all_aliases
# ---------------------------------------------------------------------------


class TestB004Rename:
    """Verify the method is now named _load_all_aliases."""

    def test_has_load_all_aliases(self):
        """EntityResolver should have _load_all_aliases method."""
        storage = AsyncMock(spec=StorageAdapter)
        llm = AsyncMock(spec=LLMProvider)
        resolver = EntityResolver(storage=storage, llm=llm)
        assert hasattr(resolver, "_load_all_aliases")
        assert callable(resolver._load_all_aliases)

    def test_no_load_aliases_paginated(self):
        """Old name _load_aliases_paginated should not exist."""
        storage = AsyncMock(spec=StorageAdapter)
        llm = AsyncMock(spec=LLMProvider)
        resolver = EntityResolver(storage=storage, llm=llm)
        assert not hasattr(resolver, "_load_aliases_paginated")

    @pytest.mark.asyncio
    async def test_load_all_aliases_returns_all(self):
        """_load_all_aliases should return all aliases."""
        aliases = [_alias("ent-1", "Alice"), _alias("ent-2", "Bob")]
        storage = _make_storage(aliases=aliases)
        llm = AsyncMock(spec=LLMProvider)
        resolver = EntityResolver(storage=storage, llm=llm)

        result = await resolver._load_all_aliases("default")
        assert len(result) == 2


# ---------------------------------------------------------------------------
# B-017: argpartition non-determinism at ties
# ---------------------------------------------------------------------------


class TestB017DeterministicOrdering:
    """Verify that tied scores produce deterministic ordering."""

    @pytest.mark.asyncio
    async def test_tied_embeddings_deterministic(self):
        """Entities with identical embeddings should produce deterministic results."""
        n = 5
        aliases = []
        entity_map = {}
        for i in range(n):
            eid = f"ent-{i}"
            name = f"Entity{i}"
            aliases.append(_alias(eid, name))
            entity_map[eid] = _entity(eid, name)

        storage = _make_storage(aliases=aliases, entity_map=entity_map)
        # All identical embeddings → all cosine similarities are 1.0 (tied)
        embeddings = [[1, 0, 0, 0]] * n
        llm = _make_llm(embeddings=embeddings)

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

        # Run multiple times and check consistency
        results = []
        for _ in range(5):
            r = await resolver.suggest_merges(owner_id="default")
            results.append([(s.keep_id, s.merge_id) for s in r])

        # All runs should produce identical results
        for r in results[1:]:
            assert r == results[0], "Non-deterministic results at ties"


# ---------------------------------------------------------------------------
# B-030: candidate_k=0 materializes full O(n²) set
# ---------------------------------------------------------------------------


class TestB030CandidateKZero:
    """Verify candidate_k=0 returns empty instead of O(n²) fallthrough."""

    @pytest.mark.asyncio
    async def test_candidate_k_zero_returns_empty(self):
        """candidate_k=0 with embedding prefilter should return [] early."""
        n = 5
        aliases = []
        entity_map = {}
        for i in range(n):
            eid = f"ent-{i}"
            # Use very similar names so fuzzy would match in O(n²) mode
            name = f"Andrew{i}"
            aliases.append(_alias(eid, name))
            entity_map[eid] = _entity(eid, name)

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

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

        result = await resolver.suggest_merges(
            owner_id="default",
            use_embedding_prefilter=True,
        )
        assert result == []
        # LLM embed should NOT have been called
        llm.embed.assert_not_called()

    @pytest.mark.asyncio
    async def test_candidate_k_negative_returns_empty(self):
        """candidate_k < 0 with embedding prefilter should return [] early."""
        aliases = [_alias("ent-1", "Alice"), _alias("ent-2", "Alicia")]
        entity_map = {
            "ent-1": _entity("ent-1", "Alice"),
            "ent-2": _entity("ent-2", "Alicia"),
        }
        storage = _make_storage(aliases=aliases, entity_map=entity_map)
        llm = _make_llm()

        resolver = EntityResolver(
            storage=storage, llm=llm, candidate_k=-1,
        )

        result = await resolver.suggest_merges(
            owner_id="default",
            use_embedding_prefilter=True,
        )
        assert result == []

    @pytest.mark.asyncio
    async def test_candidate_k_zero_without_prefilter_still_works(self):
        """candidate_k=0 with use_embedding_prefilter=False should still do O(n²) fuzzy."""
        aliases = [
            _alias("ent-1", "Andrew"),
            _alias("ent-2", "Andrew Smith"),
        ]
        entity_map = {
            "ent-1": _entity("ent-1", "Andrew"),
            "ent-2": _entity("ent-2", "Andrew Smith"),
        }
        storage = _make_storage(aliases=aliases, entity_map=entity_map)
        llm = _make_llm()

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

        result = await resolver.suggest_merges(
            owner_id="default",
            use_embedding_prefilter=False,  # Explicitly disable
        )
        # Should still find matches via O(n²) fuzzy
        assert isinstance(result, list)
