"""EpisodicMemory 向量检索单元测试 - cosine similarity + hybrid scoring"""

import uuid
from contextlib import asynccontextmanager
from datetime import datetime, timedelta, timezone
from unittest.mock import AsyncMock, MagicMock

import pytest
from sqlalchemy import Column, DateTime, Float, String
from sqlalchemy.orm import DeclarativeBase

from agentkit.memory.episodic import EpisodicMemory
from agentkit.memory.base import MemoryItem
from agentkit.memory.embedder import MockEmbedder


# ── 真实 SQLAlchemy 模型（用于测试） ─────────────────────


class Base(DeclarativeBase):
    pass


class MockEpisodicModel(Base):
    """模拟 EpisodicMemory ORM 模型"""

    __tablename__ = "test_episodic_vector_search"

    id = Column(String, primary_key=True, default=lambda: str(uuid.uuid4()))
    agent_name = Column(String, default="")
    task_type = Column(String, default="")
    input_summary = Column(String, default="")
    output_summary = Column(String, default="")
    outcome = Column(String, default="success")
    quality_score = Column(Float, default=0.5)
    reflection = Column(String, default="")
    embedding = Column(String, nullable=True)
    created_at = Column(DateTime, default=lambda: datetime.now(timezone.utc))


# ── Mock 辅助工具 ────────────────────────────────────────


def make_mock_entry(
    id: uuid.UUID | None = None,
    agent_name: str = "test_agent",
    task_type: str = "analysis",
    input_summary: str = "test input",
    output_summary: str = "test output",
    outcome: str = "success",
    quality_score: float = 0.8,
    reflection: str = "",
    embedding: list[float] | None = None,
    created_at: datetime | None = None,
):
    """创建一个模拟的 ORM entry 对象"""
    entry = MockEpisodicModel(
        id=str(id or uuid.uuid4()),
        agent_name=agent_name,
        task_type=task_type,
        input_summary=input_summary,
        output_summary=output_summary,
        outcome=outcome,
        quality_score=quality_score,
        reflection=reflection,
        created_at=created_at or datetime.now(timezone.utc),
    )
    # 直接设置 embedding 属性（绕过 Column 限制）
    entry.embedding = embedding
    return entry


def make_mock_session_factory(entries: list | None = None):
    """创建一个 mock session_factory"""
    entries = entries or []

    mock_session = AsyncMock()
    mock_session.add = MagicMock()
    mock_session.commit = AsyncMock()
    mock_session.rollback = AsyncMock()

    mock_result = MagicMock()
    mock_scalars = MagicMock()
    mock_scalars.all.return_value = entries
    mock_result.scalars.return_value = mock_scalars
    mock_session.execute = AsyncMock(return_value=mock_result)

    @asynccontextmanager
    async def factory():
        yield mock_session

    return factory, mock_session


# ── Cosine Similarity 测试 ──────────────────────────────


class TestCosineSimilarity:
    """_compute_cosine_similarity 测试"""

    def test_identical_vectors_return_one(self):
        """相同向量余弦相似度为 1"""
        vec = [1.0, 0.0, 0.0]
        assert EpisodicMemory._compute_cosine_similarity(vec, vec) == pytest.approx(1.0)

    def test_orthogonal_vectors_return_zero(self):
        """正交向量余弦相似度为 0"""
        vec_a = [1.0, 0.0]
        vec_b = [0.0, 1.0]
        assert EpisodicMemory._compute_cosine_similarity(vec_a, vec_b) == pytest.approx(0.0)

    def test_opposite_vectors_return_minus_one(self):
        """相反向量余弦相似度为 -1"""
        vec_a = [1.0, 0.0]
        vec_b = [-1.0, 0.0]
        assert EpisodicMemory._compute_cosine_similarity(vec_a, vec_b) == pytest.approx(-1.0)

    def test_dimension_mismatch_returns_zero(self):
        """维度不匹配返回 0"""
        vec_a = [1.0, 2.0]
        vec_b = [1.0]
        assert EpisodicMemory._compute_cosine_similarity(vec_a, vec_b) == 0.0

    def test_empty_vectors_return_zero(self):
        """空向量返回 0"""
        assert EpisodicMemory._compute_cosine_similarity([], []) == 0.0

    def test_zero_vector_returns_zero(self):
        """零向量返回 0"""
        vec_a = [0.0, 0.0]
        vec_b = [1.0, 2.0]
        assert EpisodicMemory._compute_cosine_similarity(vec_a, vec_b) == 0.0


# ── MockEmbedder 测试 ───────────────────────────────────


class TestMockEmbedder:
    """MockEmbedder 测试"""

    async def test_embed_returns_correct_dimension(self):
        """embed 返回指定维度的向量"""
        embedder = MockEmbedder(dimension=64)
        vec = await embedder.embed("test text")
        assert len(vec) == 64

    async def test_embed_is_deterministic(self):
        """相同文本生成相同向量"""
        embedder = MockEmbedder(dimension=32)
        vec1 = await embedder.embed("hello world")
        vec2 = await embedder.embed("hello world")
        assert vec1 == vec2

    async def test_embed_different_text_different_vector(self):
        """不同文本生成不同向量"""
        embedder = MockEmbedder(dimension=32)
        vec1 = await embedder.embed("hello")
        vec2 = await embedder.embed("world")
        assert vec1 != vec2

    async def test_embed_produces_unit_vector(self):
        """embed 生成单位向量"""
        embedder = MockEmbedder(dimension=32)
        vec = await embedder.embed("test")
        magnitude = sum(x**2 for x in vec) ** 0.5
        assert magnitude == pytest.approx(1.0, abs=1e-6)

    def test_get_dimension(self):
        """get_dimension 返回正确维度"""
        embedder = MockEmbedder(dimension=256)
        assert embedder.get_dimension() == 256


# ── Store 测试 ──────────────────────────────────────────


class TestStoreWithEmbedder:
    """store() 带 embedder 的测试"""

    async def test_store_generates_embedding_when_embedder_provided(self):
        """有 embedder 时 store 生成 embedding"""
        factory, mock_session = make_mock_session_factory()
        embedder = MockEmbedder(dimension=32)

        mem = EpisodicMemory(
            session_factory=factory,
            episodic_model=MockEpisodicModel,
            embedder=embedder,
        )

        await mem.store("key1", "some value", {"agent_name": "test"})

        entry_arg = mock_session.add.call_args[0][0]
        assert entry_arg.embedding is not None
        assert len(entry_arg.embedding) == 32

    async def test_store_no_embedding_without_embedder(self):
        """无 embedder 时 store 不生成 embedding"""
        factory, mock_session = make_mock_session_factory()

        mem = EpisodicMemory(
            session_factory=factory,
            episodic_model=MockEpisodicModel,
        )

        await mem.store("key1", "some value")

        entry_arg = mock_session.add.call_args[0][0]
        assert entry_arg.embedding is None


# ── Search 向量检索测试 ─────────────────────────────────


class TestSearchVectorSearch:
    """search() 向量检索测试"""

    async def test_search_with_embedder_uses_cosine_similarity(self):
        """有 embedder 时 search 使用 cosine similarity 排序"""
        embedder = MockEmbedder(dimension=32)

        # 生成 embedding
        vec_similar = await embedder.embed("financial analysis")
        vec_different = await embedder.embed("completely unrelated topic xyz")

        now = datetime.now(timezone.utc)
        similar_entry = make_mock_entry(
            input_summary="financial analysis report",
            quality_score=0.5,
            embedding=vec_similar,
            created_at=now,
        )
        different_entry = make_mock_entry(
            input_summary="unrelated task",
            quality_score=0.5,
            embedding=vec_different,
            created_at=now,
        )

        factory, _ = make_mock_session_factory([similar_entry, different_entry])

        mem = EpisodicMemory(
            session_factory=factory,
            episodic_model=MockEpisodicModel,
            embedder=embedder,
            alpha=1.0,  # 纯 cosine 排序
        )

        results = await mem.search("financial analysis")
        assert len(results) == 2
        # 相似条目应排在前面
        assert results[0].value["input_summary"] == "financial analysis report"

    async def test_search_fallback_to_time_decay_without_embedder(self):
        """无 embedder 时 search 回退到时间衰减排序"""
        now = datetime.now(timezone.utc)
        recent_entry = make_mock_entry(
            quality_score=0.8,
            created_at=now - timedelta(hours=1),
        )
        old_entry = make_mock_entry(
            quality_score=0.8,
            created_at=now - timedelta(hours=100),
        )

        factory, _ = make_mock_session_factory([recent_entry, old_entry])

        mem = EpisodicMemory(
            session_factory=factory,
            episodic_model=MockEpisodicModel,
        )

        results = await mem.search("test query")
        assert len(results) == 2
        # 近期条目应排在前面（纯时间衰减）
        assert results[0].score > results[1].score

    async def test_search_hybrid_scoring_formula(self):
        """混合评分公式：alpha * cosine + (1-alpha) * time_decay"""
        embedder = MockEmbedder(dimension=32)

        vec_similar = await embedder.embed("query text")
        vec_different = await embedder.embed("something else entirely")

        now = datetime.now(timezone.utc)
        # 相似条目但质量低
        similar_entry = make_mock_entry(
            quality_score=0.5,
            embedding=vec_similar,
            created_at=now,
        )
        # 不相似条目但质量高
        different_entry = make_mock_entry(
            quality_score=0.9,
            embedding=vec_different,
            created_at=now,
        )

        factory, _ = make_mock_session_factory([similar_entry, different_entry])

        # alpha=1.0 → 纯 cosine 排序
        mem = EpisodicMemory(
            session_factory=factory,
            episodic_model=MockEpisodicModel,
            embedder=embedder,
            alpha=1.0,
        )

        results = await mem.search("query text")
        # alpha=1.0 时，cosine 主导，相似条目排前面
        assert results[0].value["input_summary"] == similar_entry.input_summary

    async def test_search_alpha_zero_pure_time_decay(self):
        """alpha=0 时完全使用时间衰减排序"""
        embedder = MockEmbedder(dimension=32)

        vec_similar = await embedder.embed("query text")
        vec_different = await embedder.embed("something else")

        now = datetime.now(timezone.utc)
        # 相似但质量低
        similar_entry = make_mock_entry(
            quality_score=0.3,
            embedding=vec_similar,
            created_at=now,
        )
        # 不相似但质量高
        different_entry = make_mock_entry(
            quality_score=0.9,
            embedding=vec_different,
            created_at=now,
        )

        factory, _ = make_mock_session_factory([similar_entry, different_entry])

        mem = EpisodicMemory(
            session_factory=factory,
            episodic_model=MockEpisodicModel,
            embedder=embedder,
            alpha=0.0,  # 纯时间衰减
        )

        results = await mem.search("query text")
        # alpha=0 时，time_decay 主导，高质量条目排前面
        assert results[0].value["quality_score"] == 0.9

    async def test_search_entry_without_embedding_uses_time_decay(self):
        """有 embedder 但 entry 没有 embedding 时使用时间衰减"""
        embedder = MockEmbedder(dimension=32)

        now = datetime.now(timezone.utc)
        entry_with_embedding = make_mock_entry(
            quality_score=0.5,
            embedding=await embedder.embed("test"),
            created_at=now - timedelta(hours=10),
        )
        entry_without_embedding = make_mock_entry(
            quality_score=0.9,
            embedding=None,
            created_at=now,
        )

        factory, _ = make_mock_session_factory([entry_with_embedding, entry_without_embedding])

        mem = EpisodicMemory(
            session_factory=factory,
            episodic_model=MockEpisodicModel,
            embedder=embedder,
            alpha=0.7,
        )

        results = await mem.search("test query")
        assert len(results) == 2

    async def test_search_empty_store_returns_empty(self):
        """空存储 search 返回空列表"""
        factory, _ = make_mock_session_factory([])
        embedder = MockEmbedder(dimension=32)

        mem = EpisodicMemory(
            session_factory=factory,
            episodic_model=MockEpisodicModel,
            embedder=embedder,
        )

        results = await mem.search("anything")
        assert results == []


# ── Retrieve 向量检索测试 ───────────────────────────────


class TestRetrieveVectorSearch:
    """retrieve() 向量检索测试"""

    async def test_retrieve_with_embedder_returns_best_match(self):
        """有 embedder 时 retrieve 返回最相似条目"""
        embedder = MockEmbedder(dimension=32)

        vec_similar = await embedder.embed("financial report")
        vec_different = await embedder.embed("weather forecast")

        now = datetime.now(timezone.utc)
        similar_entry = make_mock_entry(
            input_summary="financial report Q4",
            embedding=vec_similar,
            created_at=now,
        )
        different_entry = make_mock_entry(
            input_summary="weather forecast today",
            embedding=vec_different,
            created_at=now,
        )

        factory, _ = make_mock_session_factory([similar_entry, different_entry])

        mem = EpisodicMemory(
            session_factory=factory,
            episodic_model=MockEpisodicModel,
            embedder=embedder,
        )

        result = await mem.retrieve("financial report")
        assert result is not None
        assert result.value["input_summary"] == "financial report Q4"
        assert result.metadata["cosine_similarity"] > 0.0

    async def test_retrieve_without_embedder_returns_none(self):
        """无 embedder 时 retrieve 返回 None"""
        factory, _ = make_mock_session_factory([])

        mem = EpisodicMemory(
            session_factory=factory,
            episodic_model=MockEpisodicModel,
        )

        result = await mem.retrieve("any key")
        assert result is None

    async def test_retrieve_empty_store_returns_none(self):
        """空存储 retrieve 返回 None"""
        factory, _ = make_mock_session_factory([])
        embedder = MockEmbedder(dimension=32)

        mem = EpisodicMemory(
            session_factory=factory,
            episodic_model=MockEpisodicModel,
            embedder=embedder,
        )

        result = await mem.retrieve("any key")
        assert result is None

    async def test_retrieve_no_entries_with_embedding_returns_none(self):
        """所有 entry 都没有 embedding 时 retrieve 返回 None"""
        embedder = MockEmbedder(dimension=32)

        now = datetime.now(timezone.utc)
        entry = make_mock_entry(
            embedding=None,
            created_at=now,
        )

        factory, _ = make_mock_session_factory([entry])

        mem = EpisodicMemory(
            session_factory=factory,
            episodic_model=MockEpisodicModel,
            embedder=embedder,
        )

        result = await mem.retrieve("any key")
        assert result is None

    async def test_retrieve_returns_memory_item(self):
        """retrieve 返回 MemoryItem 实例"""
        embedder = MockEmbedder(dimension=32)

        vec = await embedder.embed("test query")
        now = datetime.now(timezone.utc)
        entry = make_mock_entry(
            input_summary="test input",
            output_summary="test output",
            outcome="success",
            quality_score=0.9,
            embedding=vec,
            created_at=now,
        )

        factory, _ = make_mock_session_factory([entry])

        mem = EpisodicMemory(
            session_factory=factory,
            episodic_model=MockEpisodicModel,
            embedder=embedder,
        )

        result = await mem.retrieve("test query")
        assert isinstance(result, MemoryItem)
        assert result.value["input_summary"] == "test input"
        assert result.value["output_summary"] == "test output"
        assert result.value["outcome"] == "success"
        assert result.score > 0.0


# ── Alpha 参数测试 ──────────────────────────────────────


class TestAlphaParameter:
    """alpha 参数控制混合评分平衡"""

    async def test_alpha_controls_hybrid_balance(self):
        """alpha 控制语义相似度和时间衰减的平衡"""
        embedder = MockEmbedder(dimension=32)

        vec_similar = await embedder.embed("machine learning")
        vec_different = await embedder.embed("cooking recipes")

        now = datetime.now(timezone.utc)
        similar_entry = make_mock_entry(
            quality_score=0.3,
            embedding=vec_similar,
            created_at=now,
        )
        different_entry = make_mock_entry(
            quality_score=0.9,
            embedding=vec_different,
            created_at=now,
        )

        # alpha=1.0: 纯 cosine → 相似条目排前面
        factory1, _ = make_mock_session_factory([similar_entry, different_entry])
        mem_high_alpha = EpisodicMemory(
            session_factory=factory1,
            episodic_model=MockEpisodicModel,
            embedder=embedder,
            alpha=1.0,
        )
        results_high = await mem_high_alpha.search("machine learning")
        assert results_high[0].value["quality_score"] == 0.3  # 相似条目

        # alpha=0.0: 纯 time_decay → 高质量条目排前面
        factory2, _ = make_mock_session_factory([similar_entry, different_entry])
        mem_low_alpha = EpisodicMemory(
            session_factory=factory2,
            episodic_model=MockEpisodicModel,
            embedder=embedder,
            alpha=0.0,
        )
        results_low = await mem_low_alpha.search("machine learning")
        assert results_low[0].value["quality_score"] == 0.9  # 高质量条目

    async def test_default_alpha_is_0_7(self):
        """默认 alpha 值为 0.7"""
        factory, _ = make_mock_session_factory([])

        mem = EpisodicMemory(
            session_factory=factory,
            episodic_model=MockEpisodicModel,
        )

        assert mem._alpha == 0.7