"""Tests for semantic memory: embeddings, memory store, tool handlers, LLM client."""

import json
import math
import tempfile
import os
import pytest
from unittest.mock import patch, MagicMock, call


# ---------------------------------------------------------------------------
# 1. TestCosine
# ---------------------------------------------------------------------------

class TestCosine:
    """Tests for core.embeddings.cosine_similarity."""

    def test_identical_vectors(self):
        from core.embeddings import cosine_similarity
        v = [1.0, 2.0, 3.0]
        assert cosine_similarity(v, v) == pytest.approx(1.0)

    def test_orthogonal_vectors(self):
        from core.embeddings import cosine_similarity
        a = [1.0, 0.0]
        b = [0.0, 1.0]
        assert cosine_similarity(a, b) == pytest.approx(0.0)

    def test_opposite_vectors(self):
        from core.embeddings import cosine_similarity
        a = [1.0, 2.0, 3.0]
        b = [-1.0, -2.0, -3.0]
        assert cosine_similarity(a, b) == pytest.approx(-1.0)

    def test_dimension_mismatch_raises(self):
        from core.embeddings import cosine_similarity
        with pytest.raises(ValueError):
            cosine_similarity([1.0, 2.0], [1.0, 2.0, 3.0])

    def test_zero_vector_returns_zero(self):
        from core.embeddings import cosine_similarity
        a = [0.0, 0.0, 0.0]
        b = [1.0, 2.0, 3.0]
        assert cosine_similarity(a, b) == pytest.approx(0.0)


# ---------------------------------------------------------------------------
# 2. TestTopKSimilar
# ---------------------------------------------------------------------------

class TestTopKSimilar:
    """Tests for core.embeddings.top_k_similar."""

    def test_correct_ranking(self):
        from core.embeddings import top_k_similar
        query = [1.0, 0.0]
        candidates = [
            ("a", [1.0, 0.0]),   # identical -> 1.0
            ("b", [0.0, 1.0]),   # orthogonal -> 0.0
            ("c", [0.7, 0.7]),   # partial
        ]
        result = top_k_similar(query, candidates, k=3)
        ids = [r[0] for r in result]
        assert ids[0] == "a"
        assert ids[-1] == "b"

    def test_k_limit_respected(self):
        from core.embeddings import top_k_similar
        query = [1.0, 0.0]
        candidates = [
            ("a", [1.0, 0.0]),
            ("b", [0.5, 0.5]),
            ("c", [0.0, 1.0]),
        ]
        result = top_k_similar(query, candidates, k=2)
        assert len(result) == 2

    def test_empty_candidates(self):
        from core.embeddings import top_k_similar
        result = top_k_similar([1.0, 0.0], [], k=5)
        assert result == []


# ---------------------------------------------------------------------------
# 3. TestMemoryEntryEmbedding
# ---------------------------------------------------------------------------

class TestMemoryEntryEmbedding:
    """Tests for MemoryEntry embedding serialization."""

    def test_to_dict_includes_embedding(self):
        from core.memory_store import MemoryEntry
        entry = MemoryEntry(text="hello", tags=["t"], embedding=[0.1, 0.2])
        d = entry.to_dict()
        assert "embedding" in d
        assert d["embedding"] == [0.1, 0.2]

    def test_to_dict_excludes_embedding_when_none(self):
        from core.memory_store import MemoryEntry
        entry = MemoryEntry(text="hello", tags=["t"], embedding=None)
        d = entry.to_dict()
        assert d.get("embedding") is None

    def test_from_dict_reads_embedding(self):
        from core.memory_store import MemoryEntry
        data = {
            "text": "hello",
            "tags": ["t"],
            "entry_id": "abc",
            "source": "test",
            "created_at": "2025-01-01T00:00:00",
            "updated_at": "2025-01-01T00:00:00",
            "embedding": [0.3, 0.4],
        }
        entry = MemoryEntry.from_dict(data)
        assert entry.embedding == [0.3, 0.4]

    def test_from_dict_backward_compat_no_embedding(self):
        from core.memory_store import MemoryEntry
        data = {
            "text": "hello",
            "tags": ["t"],
            "entry_id": "abc",
            "source": "test",
            "created_at": "2025-01-01T00:00:00",
            "updated_at": "2025-01-01T00:00:00",
        }
        entry = MemoryEntry.from_dict(data)
        assert entry.embedding is None


# ---------------------------------------------------------------------------
# 4. TestMemoryStoreSemanticRecall
# ---------------------------------------------------------------------------

class TestMemoryStoreSemanticRecall:
    """Tests for MemoryStore.semantic_recall."""

    def _make_store(self, tmp):
        from core.memory_store import MemoryStore
        return MemoryStore(store_dir=tmp)

    def test_returns_correct_order(self):
        from core.memory_store import MemoryStore
        with tempfile.TemporaryDirectory() as tmp:
            store = self._make_store(tmp)
            store.remember("u1", "close", tags=[], source="test", embedding=[1.0, 0.0])
            store.remember("u1", "far", tags=[], source="test", embedding=[0.0, 1.0])
            store.remember("u1", "mid", tags=[], source="test", embedding=[0.7, 0.7])
            results = store.semantic_recall("u1", query_embedding=[1.0, 0.0], limit=10)
            texts = [entry.text for entry, score in results]
            assert texts[0] == "close"

    def test_limit_is_respected(self):
        with tempfile.TemporaryDirectory() as tmp:
            store = self._make_store(tmp)
            for i in range(5):
                emb = [0.0] * 3
                emb[i % 3] = 1.0
                store.remember("u1", f"entry{i}", tags=[], source="test", embedding=emb)
            results = store.semantic_recall("u1", query_embedding=[1.0, 0.0, 0.0], limit=2)
            assert len(results) <= 2

    def test_entries_without_embeddings_skipped(self):
        with tempfile.TemporaryDirectory() as tmp:
            store = self._make_store(tmp)
            store.remember("u1", "has_emb", tags=[], source="test", embedding=[1.0, 0.0])
            store.remember("u1", "no_emb", tags=[], source="test", embedding=None)
            results = store.semantic_recall("u1", query_embedding=[1.0, 0.0], limit=10)
            texts = [entry.text for entry, _ in results]
            assert "no_emb" not in texts
            assert "has_emb" in texts

    def test_empty_store_returns_empty(self):
        with tempfile.TemporaryDirectory() as tmp:
            store = self._make_store(tmp)
            results = store.semantic_recall("u1", query_embedding=[1.0, 0.0], limit=10)
            assert results == []


# ---------------------------------------------------------------------------
# 5. TestMemoryStoreRememberWithEmbedding
# ---------------------------------------------------------------------------

class TestMemoryStoreRememberWithEmbedding:
    """Tests for MemoryStore.remember with embedding parameter."""

    def test_remember_stores_embedding(self):
        from core.memory_store import MemoryStore
        with tempfile.TemporaryDirectory() as tmp:
            store = MemoryStore(store_dir=tmp)
            entry = store.remember("u1", "test text", tags=["a"], source="test", embedding=[0.5, 0.5])
            assert entry.embedding == [0.5, 0.5]

    def test_remember_updates_embedding_on_duplicate(self):
        from core.memory_store import MemoryStore
        with tempfile.TemporaryDirectory() as tmp:
            store = MemoryStore(store_dir=tmp)
            e1 = store.remember("u1", "same text", tags=["a"], source="test", embedding=[0.1, 0.2])
            e2 = store.remember("u1", "same text", tags=["a"], source="test", embedding=[0.9, 0.8])
            assert e2.embedding == [0.9, 0.8]


# ---------------------------------------------------------------------------
# 6. TestMemoryStoreReEmbedAll
# ---------------------------------------------------------------------------

class TestMemoryStoreReEmbedAll:
    """Tests for MemoryStore.re_embed_all."""

    def test_calls_embed_fn_for_each_entry(self):
        from core.memory_store import MemoryStore
        with tempfile.TemporaryDirectory() as tmp:
            store = MemoryStore(store_dir=tmp)
            store.remember("u1", "aaa", tags=[], source="test")
            store.remember("u1", "bbb", tags=[], source="test")
            embed_fn = MagicMock(return_value=[0.1, 0.2])
            store.re_embed_all("u1", embed_fn)
            assert embed_fn.call_count == 2

    def test_handles_errors_gracefully(self):
        from core.memory_store import MemoryStore
        with tempfile.TemporaryDirectory() as tmp:
            store = MemoryStore(store_dir=tmp)
            store.remember("u1", "aaa", tags=[], source="test")
            embed_fn = MagicMock(side_effect=RuntimeError("fail"))
            # Should not raise
            count = store.re_embed_all("u1", embed_fn)
            assert count == 0

    def test_returns_count(self):
        from core.memory_store import MemoryStore
        with tempfile.TemporaryDirectory() as tmp:
            store = MemoryStore(store_dir=tmp)
            store.remember("u1", "aaa", tags=[], source="test")
            store.remember("u1", "bbb", tags=[], source="test")
            store.remember("u1", "ccc", tags=[], source="test")
            embed_fn = MagicMock(return_value=[0.1, 0.2])
            count = store.re_embed_all("u1", embed_fn)
            assert count == 3


# ---------------------------------------------------------------------------
# 7. TestSemanticRecallHandler
# ---------------------------------------------------------------------------

class TestSemanticRecallHandler:
    """Tests for SemanticRecallHandler from core.tool_registry."""

    def test_execute_with_mock_embed_fn(self):
        from core.tool_registry import SemanticRecallHandler
        from core.memory_store import MemoryStore
        with tempfile.TemporaryDirectory() as tmp:
            store = MemoryStore(store_dir=tmp)
            store.remember("u1", "relevant entry", tags=[], source="test", embedding=[1.0, 0.0])
            store.remember("u1", "irrelevant", tags=[], source="test", embedding=[0.0, 1.0])

            handler = SemanticRecallHandler()
            handler.set_user_id("u1")
            handler.set_embed_fn(lambda text: [1.0, 0.0])

            mock_instance = MagicMock()
            mock_instance.semantic_recall.return_value = store.semantic_recall("u1", [1.0, 0.0], limit=5)
            with patch("core.memory_store.MemoryStore.instance", return_value=mock_instance):
                # Use the real store directly
                pass
            # Just call directly — MemoryStore.instance() is called inside execute
            # We need to set the singleton
            old_instance = MemoryStore._instance
            MemoryStore._instance = store
            try:
                result = handler.execute({"query": "find relevant", "limit": 5})
            finally:
                MemoryStore._instance = old_instance
            assert "relevant entry" in str(result)

    def test_execute_without_embed_fn_returns_error(self):
        from core.tool_registry import SemanticRecallHandler
        handler = SemanticRecallHandler()
        handler.set_user_id("u1")
        # No embed_fn set
        result = handler.execute({"query": "test", "limit": 5})
        assert "error" in str(result).lower() or "embed" in str(result).lower()

    def test_execute_with_empty_query_returns_error(self):
        from core.tool_registry import SemanticRecallHandler
        handler = SemanticRecallHandler()
        handler.set_user_id("u1")
        handler.set_embed_fn(lambda text: [1.0, 0.0])
        result = handler.execute({"query": "", "limit": 5})
        assert "error" in str(result).lower() or result is not None


# ---------------------------------------------------------------------------
# 8. TestRememberHandlerAutoEmbed
# ---------------------------------------------------------------------------

class TestRememberHandlerAutoEmbed:
    """Tests for RememberHandler auto-embedding via embed_fn."""

    def test_embedding_stored_when_embed_fn_set(self):
        from core.tool_registry import RememberHandler
        from core.memory_store import MemoryStore
        with tempfile.TemporaryDirectory() as tmp:
            store = MemoryStore(store_dir=tmp)
            handler = RememberHandler()
            handler.set_user_id("u1")
            handler.set_embed_fn(lambda text: [0.5, 0.5])

            old = MemoryStore._instance
            MemoryStore._instance = store
            try:
                handler.execute({"text": "remember this", "tags": ["test"]})
            finally:
                MemoryStore._instance = old

            results = store.semantic_recall("u1", query_embedding=[0.5, 0.5], limit=10)
            assert len(results) >= 1

    def test_embedding_not_stored_when_no_embed_fn(self):
        from core.tool_registry import RememberHandler
        from core.memory_store import MemoryStore
        with tempfile.TemporaryDirectory() as tmp:
            store = MemoryStore(store_dir=tmp)
            handler = RememberHandler()
            handler.set_user_id("u1")

            old = MemoryStore._instance
            MemoryStore._instance = store
            try:
                handler.execute({"text": "remember this", "tags": ["test"]})
            finally:
                MemoryStore._instance = old

            results = store.semantic_recall("u1", query_embedding=[0.5, 0.5], limit=10)
            assert len(results) == 0

    def test_embed_fn_error_doesnt_prevent_remember(self):
        from core.tool_registry import RememberHandler
        from core.memory_store import MemoryStore
        with tempfile.TemporaryDirectory() as tmp:
            store = MemoryStore(store_dir=tmp)
            handler = RememberHandler()
            handler.set_user_id("u1")
            handler.set_embed_fn(MagicMock(side_effect=RuntimeError("embed fail")))

            old = MemoryStore._instance
            MemoryStore._instance = store
            try:
                result = handler.execute({"text": "remember this", "tags": ["test"]})
            finally:
                MemoryStore._instance = old
            assert "remembered" in str(result).lower()


# ---------------------------------------------------------------------------
# 9. TestLLMClientEmbed
# ---------------------------------------------------------------------------

class TestLLMClientEmbed:
    """Tests for LLMClient.embed (OpenAI embeddings)."""

    def test_mock_http_response_openai(self):
        from core.llm_client import LLMClient
        client = LLMClient(provider="openai", config={"api_key": "test-key"})
        mock_response = {
            "data": [
                {"embedding": [0.1, 0.2, 0.3]},
            ],
            "usage": {"total_tokens": 5},
        }
        with patch.object(client, '_http_post', return_value=mock_response):
            result = client.embed(["hello world"], model="text-embedding-3-small")
        assert len(result) == 1
        assert result[0] == [0.1, 0.2, 0.3]

    def test_error_for_non_openai_provider(self):
        from core.llm_client import LLMClient
        client = LLMClient(provider="anthropic", config={"api_key": "test-key"})
        with pytest.raises((ValueError, NotImplementedError, Exception)):
            client.embed(["hello"])

    def test_batching_large_input(self):
        from core.llm_client import LLMClient
        client = LLMClient(provider="openai", config={"api_key": "test-key"})

        def mock_post(path, body, headers=None):
            texts = body.get("input", [])
            return {
                "data": [{"embedding": [0.1, 0.2], "index": i} for i, _ in enumerate(texts)],
                "usage": {"total_tokens": len(texts)},
            }

        with patch.object(client, '_http_post', side_effect=mock_post):
            texts = [f"text {i}" for i in range(50)]
            result = client.embed(texts, model="text-embedding-3-small")
        assert len(result) == 50

    def test_error_missing_api_key(self):
        from core.llm_client import LLMClient
        with pytest.raises(Exception):
            client = LLMClient(provider="openai", config={"api_key": ""})
            client.embed(["hello"])


# ---------------------------------------------------------------------------
# 10. TestEmbeddingProvider
# ---------------------------------------------------------------------------

class TestEmbeddingProvider:
    """Tests for EmbeddingProvider auto-detection."""

    def test_auto_detection_with_api_key_returns_openai(self):
        from core.embeddings import EmbeddingProvider
        provider = EmbeddingProvider.instance()
        # With an api_key, should select openai
        with patch.object(provider, 'embed', wraps=provider.embed) as mock_embed:
            try:
                provider.embed(["test"], provider="auto", api_key="sk-test123")
            except Exception:
                pass
            # Check that the method was called; the provider resolution should pick openai
            # We verify by checking it doesn't raise ValueError for "auto"

    def test_auto_detection_without_api_key_returns_local(self, monkeypatch):
        # Forbid HuggingFace downloads: on a fresh CI runner the local
        # provider would fetch a real model, and a stalled hf_xet download
        # hangs the whole suite until the 6h Actions limit. With a cached
        # model the embed still works offline; without one it raises fast
        # and the except branch below absorbs it.
        monkeypatch.setenv("HF_HUB_OFFLINE", "1")
        from core.embeddings import EmbeddingProvider
        provider = EmbeddingProvider.instance()
        try:
            result = provider.embed(["test"], provider="auto", api_key=None)
            # If local provider is available, it should return embeddings
            assert isinstance(result, list)
        except Exception:
            # Local provider may not be installed; that's OK
            pass

    def test_unknown_provider_raises(self):
        from core.embeddings import EmbeddingProvider
        provider = EmbeddingProvider.instance()
        with pytest.raises((ValueError, KeyError, Exception)):
            provider.embed(["test"], provider="nonexistent_provider_xyz")

    def test_memory_embed_fn_uses_embedding_llm_service(self, monkeypatch):
        from core.embeddings import EmbeddingProvider, build_memory_embed_fn
        from core.service_registry import ServiceRegistry

        class _Svc:
            def embed(self, texts):
                assert texts == ["hello"]
                return [[0.1, 0.2]]

        class _Reg:
            def resolve(self, service_id, user_id="", conv_id=""):
                assert service_id == "embedder"
                assert user_id == "u1"
                assert conv_id == "c1"
                return _Svc()

        monkeypatch.setattr(
            "core.expression.resolve_value",
            lambda value, **kwargs: "embedder",
        )
        monkeypatch.setattr(
            ServiceRegistry, "get_instance", staticmethod(lambda: _Reg()))
        monkeypatch.setattr(
            EmbeddingProvider, "instance",
            staticmethod(lambda: pytest.fail("local embeddings should not load")),
        )

        assert build_memory_embed_fn("u1", "c1")("hello") == [0.1, 0.2]

    def test_memory_embed_fn_falls_back_local_without_embedding_service(self, monkeypatch):
        from core.embeddings import EmbeddingProvider, build_memory_embed_fn

        class _Local:
            def embed(self, texts, provider="auto", **kwargs):
                assert provider == "local"
                assert texts == ["hello"]
                return [[0.3, 0.4]]

        monkeypatch.setattr(
            "core.expression.resolve_value",
            lambda value, **kwargs: "${embedding_llm_service}",
        )
        monkeypatch.setattr(
            EmbeddingProvider, "instance", staticmethod(lambda: _Local()))

        assert build_memory_embed_fn("u1", "c1")("hello") == [0.3, 0.4]


def test_llm_connection_service_embed_uses_embedding_model(monkeypatch):
    from services.llm_connection import LLMConnectionService

    svc = LLMConnectionService({
        "provider": "openai",
        "api_key": "sk-test",
        "embedding_model": "text-embedding-3-large",
    })
    seen = {}

    def _embed(texts, model=None):
        seen["texts"] = texts
        seen["model"] = model
        return [[1.0, 2.0]]

    monkeypatch.setattr(svc._client, "embed", _embed)

    assert svc.embed(["hello"]) == [[1.0, 2.0]]
    assert seen == {
        "texts": ["hello"],
        "model": "text-embedding-3-large",
    }
