"""Tests for memory-aware (informed) extraction."""

from __future__ import annotations

import json
from unittest.mock import AsyncMock, MagicMock

import pytest
from conftest import FakeEmbeddingProvider, FakeLLMProvider

from arandu.config import MemoryConfig
from arandu.constants import ExtractedFact
from arandu.write.informed_extract import (
    alias_lookup,
    build_extraction_context,
    informed_extract_from_message,
)

# ---------------------------------------------------------------------------
# alias_lookup
# ---------------------------------------------------------------------------


class TestAliasLookup:
    """Tests for alias_lookup — entity detection via string matching."""

    @pytest.mark.asyncio
    async def test_finds_known_entities(self):
        session = AsyncMock()
        session.execute.return_value = MagicMock(
            all=lambda: [
                ("caroline", "person:caroline"),
                ("mel", "person:mel"),
            ]
        )

        result = await alias_lookup(session, "test-user", "Hey Mel! How are you?")

        assert "person:mel" in result
        assert "person:caroline" not in result

    @pytest.mark.asyncio
    async def test_skips_short_aliases(self):
        session = AsyncMock()
        session.execute.return_value = MagicMock(
            all=lambda: [
                ("AI", "concept:ai"),
                ("mel", "person:mel"),
            ]
        )

        result = await alias_lookup(session, "test-user", "AI told Mel something")

        assert "concept:ai" not in result
        assert "person:mel" in result

    @pytest.mark.asyncio
    async def test_word_boundary_matching(self):
        session = AsyncMock()
        session.execute.return_value = MagicMock(
            all=lambda: [
                ("art", "concept:art"),
            ]
        )

        # "art" should not match inside "started"
        result = await alias_lookup(session, "test-user", "She started painting")
        assert "concept:art" not in result

        # "art" should match as standalone word
        result = await alias_lookup(session, "test-user", "She loves art")
        assert "concept:art" in result

    @pytest.mark.asyncio
    async def test_empty_aliases_returns_empty(self):
        session = AsyncMock()
        session.execute.return_value = MagicMock(all=lambda: [])

        result = await alias_lookup(session, "test-user", "Hello world")
        assert result == {}

    @pytest.mark.asyncio
    async def test_case_insensitive(self):
        session = AsyncMock()
        session.execute.return_value = MagicMock(
            all=lambda: [
                ("caroline", "person:caroline"),
            ]
        )

        result = await alias_lookup(session, "test-user", "CAROLINE went to the park")
        assert "person:caroline" in result


# ---------------------------------------------------------------------------
# build_extraction_context
# ---------------------------------------------------------------------------


class TestBuildExtractionContext:
    """Tests for context builder — pgvector-based fact retrieval."""

    @pytest.mark.asyncio
    async def test_returns_empty_for_no_entities(self):
        session = AsyncMock()
        config = MemoryConfig()

        result = await build_extraction_context(session, "test-user", [], [0.0] * 1536, config)
        assert result == ""

    @pytest.mark.asyncio
    async def test_returns_empty_for_no_embedding(self):
        session = AsyncMock()
        config = MemoryConfig()

        result = await build_extraction_context(session, "test-user", ["person:mel"], [], config)
        assert result == ""

    @pytest.mark.asyncio
    async def test_formats_context_with_facts(self):
        session = AsyncMock()
        session.execute.return_value = MagicMock(
            all=lambda: [
                ("Mel is friends with Caroline", 0.8),
                ("Mel has two kids", 0.7),
            ]
        )
        config = MemoryConfig()

        result = await build_extraction_context(session, "test-user", ["person:mel"], [0.0] * 1536, config)

        assert "EXISTING KNOWLEDGE" in result
        assert "Mel is friends with Caroline" in result
        assert "Mel has two kids" in result

    @pytest.mark.asyncio
    async def test_skips_low_similarity_facts(self):
        session = AsyncMock()
        session.execute.return_value = MagicMock(
            all=lambda: [
                ("Mel is friends with Caroline", 0.1),  # below 0.3 threshold
            ]
        )
        config = MemoryConfig()

        result = await build_extraction_context(session, "test-user", ["person:mel"], [0.0] * 1536, config)
        assert result == ""


# ---------------------------------------------------------------------------
# informed_extract_from_message
# ---------------------------------------------------------------------------


class TestInformedExtraction:
    """Tests for the full informed extraction pipeline."""

    @pytest.mark.asyncio
    async def test_cold_start_works(self):
        """With no aliases in DB, informed extraction still works (empty context)."""
        session = AsyncMock()
        # alias_lookup returns empty
        session.execute.return_value = MagicMock(all=lambda: [])

        llm_response = json.dumps(
            {
                "entities": [{"name": "Caroline", "type": "person", "aliases": []}],
                "facts": [
                    {
                        "subject": "Caroline",
                        "fact": "Caroline is a data scientist",
                        "confidence": 0.95,
                        "action": "NEW",
                        "importance_category": "stable_preference",
                    }
                ],
                "relations": [],
            }
        )
        llm = FakeLLMProvider([llm_response])
        embeddings = FakeEmbeddingProvider()
        config = MemoryConfig()

        output, meta = await informed_extract_from_message(
            session,
            "test-user",
            "I'm a data scientist",
            llm=llm,
            embeddings=embeddings,
            speaker_name="Caroline",
            config=config,
        )

        assert meta["mode"] == "informed"
        assert meta["context_entities"] == 0
        assert len(output.entities) == 1
        assert len(output.facts) == 1
        assert output.facts[0].action == "NEW"
        assert output.facts[0].importance_category == "stable_preference"

    @pytest.mark.asyncio
    async def test_with_existing_context(self):
        """With known entities, pre-retrieval injects context."""
        session = AsyncMock()

        # First call: alias_lookup
        alias_result = MagicMock(all=lambda: [("mel", "person:mel")])
        # Second call: load_entity_profiles (no profiles yet)
        profiles_result = MagicMock(all=lambda: [])
        # Third call: build_extraction_context (pgvector search)
        facts_result = MagicMock(all=lambda: [("Mel has two kids", 0.8)])
        session.execute.side_effect = [alias_result, profiles_result, facts_result]

        llm_response = json.dumps(
            {
                "entities": [{"name": "Mel", "type": "person", "aliases": []}],
                "facts": [
                    {
                        "subject": "Mel",
                        "fact": "Mel's youngest took her first steps",
                        "confidence": 0.95,
                        "action": "NEW",
                        "importance_category": "biographical_milestone",
                    }
                ],
                "relations": [],
            }
        )
        llm = FakeLLMProvider([llm_response])
        embeddings = FakeEmbeddingProvider()
        config = MemoryConfig()

        output, meta = await informed_extract_from_message(
            session,
            "test-user",
            "Mel's youngest just took her first steps!",
            llm=llm,
            embeddings=embeddings,
            speaker_name="Caroline",
            config=config,
        )

        assert meta["mode"] == "informed"
        assert meta["context_entities"] == 1
        assert "person:mel" in meta["context_entity_keys"]
        assert len(output.facts) == 1
        assert output.facts[0].importance_category == "biographical_milestone"

    @pytest.mark.asyncio
    async def test_llm_failure_raises(self):
        """If LLM returns None, informed extraction raises (pipeline catches)."""
        session = AsyncMock()
        session.execute.return_value = MagicMock(all=lambda: [])

        llm = FakeLLMProvider([])  # will return default empty extraction
        # Override to return None
        llm._responses = []

        embeddings = FakeEmbeddingProvider()
        config = MemoryConfig()

        # FakeLLMProvider returns '{"entities":[],...}' by default, not None
        # So this should succeed — but with an LLM that returns None:
        class NullLLM:
            async def complete(self, **kwargs):
                from arandu.protocols import LLMResult

                return LLMResult(text=None)

        # _call_llm returns None for None text, which triggers the raise
        with pytest.raises(ValueError):
            await informed_extract_from_message(
                session,
                "test-user",
                "Hello",
                llm=NullLLM(),
                embeddings=embeddings,
                speaker_name="Test",
                config=config,
            )

    @pytest.mark.asyncio
    async def test_extracted_fact_defaults(self):
        """ExtractedFact has correct defaults for backward compatibility."""
        fact = ExtractedFact(subject="X", fact="X likes coffee")
        assert fact.action == "NEW"
        assert fact.importance_category == "specific_event"
        assert fact.updates is None


# ---------------------------------------------------------------------------
# Pipeline fallback integration
# ---------------------------------------------------------------------------


class TestPipelineFallback:
    """Test that the pipeline falls back to blind extraction on failure."""

    @pytest.mark.asyncio
    async def test_fallback_on_informed_failure(self):
        """Pipeline should use blind extraction when informed extraction fails."""
        # Informed extraction failure causes fallback to blind extraction.
        # In both cases, reconcile_facts() runs unconditionally.
        meta = {}  # empty meta = informed failed, blind used
        assert meta.get("mode") != "informed"

        meta = {"mode": "informed"}
        assert meta.get("mode") == "informed"

    @pytest.mark.asyncio
    async def test_config_disables_informed(self):
        """enable_informed_extraction=False skips informed extraction entirely."""
        config = MemoryConfig(enable_informed_extraction=False)
        assert not config.enable_informed_extraction

        config = MemoryConfig(enable_informed_extraction=True)
        assert config.enable_informed_extraction


# ---------------------------------------------------------------------------
# Entity profile tests
# ---------------------------------------------------------------------------


class TestEntityProfiles:
    """Tests for living entity profile integration."""

    @pytest.mark.asyncio
    async def test_load_entity_profiles_returns_profiles(self):
        """load_entity_profiles returns dict of entity_key → profile_text."""
        from arandu.write.informed_extract import load_entity_profiles

        session = AsyncMock()
        session.execute.return_value = MagicMock(
            all=lambda: [
                ("person:mel", "Mel is a close friend of Caroline with two kids."),
            ]
        )

        profiles = await load_entity_profiles(session, "test-user", ["person:mel"])
        assert profiles == {"person:mel": "Mel is a close friend of Caroline with two kids."}

    @pytest.mark.asyncio
    async def test_load_entity_profiles_empty(self):
        """With no entity keys, returns empty dict."""
        from arandu.write.informed_extract import load_entity_profiles

        session = AsyncMock()
        profiles = await load_entity_profiles(session, "test-user", [])
        assert profiles == {}

    @pytest.mark.asyncio
    async def test_profile_used_as_context_instead_of_facts(self):
        """When profile exists, build_extraction_context uses it instead of fact retrieval."""
        session = AsyncMock()
        # Should NOT call execute for fact retrieval when profile exists
        config = MemoryConfig()

        result = await build_extraction_context(
            session,
            "test-user",
            ["person:mel"],
            [0.0] * 1536,
            config,
            entity_profiles={"person:mel": "Mel is a parent of two."},
        )

        assert "EXISTING KNOWLEDGE" in result
        assert "Profile of Mel" in result
        assert "Mel is a parent of two" in result
        # session.execute should NOT have been called (profile skips fact retrieval)
        session.execute.assert_not_called()

    @pytest.mark.asyncio
    async def test_fallback_to_facts_when_no_profile(self):
        """When entity has no profile, falls back to fact retrieval."""
        session = AsyncMock()
        session.execute.return_value = MagicMock(all=lambda: [("Mel has two kids", 0.8)])
        config = MemoryConfig()

        result = await build_extraction_context(
            session,
            "test-user",
            ["person:mel"],
            [0.0] * 1536,
            config,
            entity_profiles={},  # no profile for mel
        )

        assert "Known facts about Mel" in result
        # session.execute WAS called for fact retrieval
        session.execute.assert_called_once()

    @pytest.mark.asyncio
    async def test_updated_profiles_parsed_from_output(self):
        """informed_extract_from_message returns updated_profiles in meta."""
        session = AsyncMock()
        session.execute.return_value = MagicMock(all=lambda: [])

        llm_response = json.dumps(
            {
                "entities": [{"name": "Caroline", "type": "person", "aliases": []}],
                "facts": [
                    {
                        "subject": "Caroline",
                        "fact": "Caroline is a data scientist",
                        "confidence": 0.95,
                        "action": "NEW",
                        "importance_category": "stable_preference",
                    }
                ],
                "relations": [],
                "updated_profiles": {
                    "Caroline": "Caroline is a data scientist and a close friend of Mel.",
                },
            }
        )
        llm = FakeLLMProvider([llm_response])
        embeddings = FakeEmbeddingProvider()
        config = MemoryConfig()

        output, meta = await informed_extract_from_message(
            session,
            "test-user",
            "I'm a data scientist",
            llm=llm,
            embeddings=embeddings,
            speaker_name="Caroline",
            config=config,
        )

        assert "updated_profiles" in meta
        assert meta["updated_profiles"]["Caroline"] == "Caroline is a data scientist and a close friend of Mel."

    @pytest.mark.asyncio
    async def test_missing_updated_profiles_returns_empty(self):
        """If LLM omits updated_profiles, meta has empty dict."""
        session = AsyncMock()
        session.execute.return_value = MagicMock(all=lambda: [])

        llm_response = json.dumps(
            {
                "entities": [{"name": "Caroline", "type": "person", "aliases": []}],
                "facts": [],
                "relations": [],
                # no updated_profiles key
            }
        )
        llm = FakeLLMProvider([llm_response])
        embeddings = FakeEmbeddingProvider()
        config = MemoryConfig()

        output, meta = await informed_extract_from_message(
            session,
            "test-user",
            "Hello",
            llm=llm,
            embeddings=embeddings,
            speaker_name="Caroline",
            config=config,
        )

        assert meta["updated_profiles"] == {}

    @pytest.mark.asyncio
    async def test_null_profile_preserved(self):
        """null profile value means 'no change' — preserved in meta."""
        session = AsyncMock()
        session.execute.return_value = MagicMock(all=lambda: [])

        llm_response = json.dumps(
            {
                "entities": [{"name": "Mel", "type": "person", "aliases": []}],
                "facts": [],
                "relations": [],
                "updated_profiles": {"Mel": None},
            }
        )
        llm = FakeLLMProvider([llm_response])
        embeddings = FakeEmbeddingProvider()
        config = MemoryConfig()

        output, meta = await informed_extract_from_message(
            session,
            "test-user",
            "Hey Mel",
            llm=llm,
            embeddings=embeddings,
            speaker_name="Caroline",
            config=config,
        )

        assert meta["updated_profiles"]["Mel"] is None
