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

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
from agentkit.utils.vector_math import compute_cosine_similarity


# ── 真实 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


class _RowMapping(dict):
    """A dict subclass that supports both ``row["key"]`` and ``row.get("key")``
    access patterns, mimicking SQLAlchemy's MappingResult rows."""

    def __getattr__(self, name: str):
        try:
            return self[name]
        except KeyError:
            raise AttributeError(name)


def _make_row_mapping(data: dict) -> _RowMapping:
    """Create a _RowMapping from a dict, for use in pgvector mock tests."""
    return _RowMapping(data)


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


class TestCosineSimilarity:
    """compute_cosine_similarity 测试"""

    def test_identical_vectors_return_one(self):
        """相同向量余弦相似度为 1"""
        vec = [1.0, 0.0, 0.0]
        assert 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 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 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 compute_cosine_similarity(vec_a, vec_b) == 0.0

    def test_empty_vectors_return_zero(self):
        """空向量返回 0"""
        assert 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 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 排序
            pgvector_enabled=False,  # 使用客户端 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,
            pgvector_enabled=False,
        )

        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,  # 纯时间衰减
            pgvector_enabled=False,
        )

        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,
            pgvector_enabled=False,
        )

        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,
            pgvector_enabled=False,
        )

        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,
            pgvector_enabled=False,
        )

        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,
            pgvector_enabled=False,
        )

        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,
            pgvector_enabled=False,
        )
        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,
            pgvector_enabled=False,
        )
        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


# ── pgvector 参数测试 ───────────────────────────────────


class TestPgvectorParameters:
    """pgvector_enabled 和 table_name 参数测试"""

    def test_default_pgvector_enabled_is_true(self):
        """默认 pgvector_enabled 为 True"""
        factory, _ = make_mock_session_factory()

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

        assert mem._pgvector_enabled is True

    def test_pgvector_enabled_can_be_disabled(self):
        """可以禁用 pgvector"""
        factory, _ = make_mock_session_factory()

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

        assert mem._pgvector_enabled is False

    def test_default_table_name(self):
        """默认 table_name 为 episodic_memories"""
        factory, _ = make_mock_session_factory()

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

        assert mem._table_name == "episodic_memories"

    def test_custom_table_name(self):
        """可以自定义 table_name"""
        factory, _ = make_mock_session_factory()

        mem = EpisodicMemory(
            session_factory=factory,
            episodic_model=MockEpisodicModel,
            table_name="custom_memories",
        )

        assert mem._table_name == "custom_memories"

    async def test_search_uses_client_side_when_pgvector_disabled(self):
        """pgvector_enabled=False 时使用客户端 cosine similarity"""
        embedder = MockEmbedder(dimension=32)

        vec_similar = await embedder.embed("test query")
        vec_different = await embedder.embed("unrelated")

        now = datetime.now(timezone.utc)
        similar_entry = make_mock_entry(
            input_summary="similar task",
            quality_score=0.5,
            embedding=vec_similar,
            created_at=now,
        )
        different_entry = make_mock_entry(
            input_summary="different task",
            quality_score=0.5,
            embedding=vec_different,
            created_at=now,
        )

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

        mem = EpisodicMemory(
            session_factory=factory,
            episodic_model=MockEpisodicModel,
            embedder=embedder,
            alpha=1.0,
            pgvector_enabled=False,
        )

        results = await mem.search("test query")
        assert len(results) == 2
        # Client-side should still rank similar entry first
        assert results[0].value["input_summary"] == "similar task"

    async def test_search_uses_client_side_when_no_embedder(self):
        """没有 embedder 时即使 pgvector_enabled=True 也使用客户端路径"""
        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,
            pgvector_enabled=True,  # Enabled but no embedder → falls back
        )

        results = await mem.search("test query")
        assert len(results) == 2
        assert results[0].score > results[1].score

    async def test_retrieve_uses_client_side_when_pgvector_disabled(self):
        """pgvector_enabled=False 时 retrieve 使用客户端 cosine similarity"""
        embedder = MockEmbedder(dimension=32)

        vec = await embedder.embed("test query")
        now = datetime.now(timezone.utc)
        entry = make_mock_entry(
            input_summary="test input",
            embedding=vec,
            created_at=now,
        )

        factory, _ = make_mock_session_factory([entry])

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

        result = await mem.retrieve("test query")
        assert result is not None
        assert result.value["input_summary"] == "test input"


# ── pgvector 原生查询 Mock 测试 ─────────────────────────


class TestPgvectorNativeSearch:
    """pgvector 原生 ``<=>`` 算符检索测试（使用 mock session）"""

    async def test_search_pgvector_uses_text_query(self):
        """pgvector search 使用 SQLAlchemy text() 查询"""
        embedder = MockEmbedder(dimension=32)
        vec = await embedder.embed("test query")

        now = datetime.now(timezone.utc)

        # Mock the pgvector raw query result as a dict-like MappingRow
        mock_row = _make_row_mapping({
            "id": str(uuid.uuid4()),
            "agent_name": "test_agent",
            "task_type": "analysis",
            "input_summary": "test input",
            "output_summary": "test output",
            "outcome": "success",
            "quality_score": 0.8,
            "reflection": "",
            "embedding": vec,
            "created_at": now,
            "distance": 0.1,
        })

        mock_result = MagicMock()
        mock_result.mappings.return_value.all.return_value = [mock_row]

        mock_session = AsyncMock()
        mock_session.execute = AsyncMock(return_value=mock_result)

        @asynccontextmanager
        async def factory():
            yield mock_session

        mem = EpisodicMemory(
            session_factory=factory,
            episodic_model=MockEpisodicModel,
            embedder=embedder,
            pgvector_enabled=True,
            table_name="episodic_memories",
        )

        results = await mem.search("test query")
        assert len(results) == 1
        assert results[0].value["input_summary"] == "test input"

        # Verify that execute was called with a text() query
        mock_session.execute.assert_called_once()
        call_args = mock_session.execute.call_args
        sql_obj = call_args[0][0]
        # The SQL should contain the <=> operator
        assert "<=>" in str(sql_obj)

    async def test_retrieve_pgvector_uses_text_query(self):
        """pgvector retrieve 使用 SQLAlchemy text() 查询"""
        embedder = MockEmbedder(dimension=32)
        vec = await embedder.embed("test query")

        now = datetime.now(timezone.utc)

        mock_row = _make_row_mapping({
            "id": str(uuid.uuid4()),
            "agent_name": "test_agent",
            "task_type": "analysis",
            "input_summary": "test input",
            "output_summary": "test output",
            "outcome": "success",
            "quality_score": 0.8,
            "reflection": "",
            "embedding": vec,
            "created_at": now,
        })

        mock_result = MagicMock()
        mock_result.mappings.return_value.first.return_value = mock_row

        mock_session = AsyncMock()
        mock_session.execute = AsyncMock(return_value=mock_result)

        @asynccontextmanager
        async def factory():
            yield mock_session

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

        result = await mem.retrieve("test query")
        assert result is not None
        assert result.value["input_summary"] == "test input"

        # Verify that execute was called with a text() query
        mock_session.execute.assert_called_once()
        call_args = mock_session.execute.call_args
        sql_obj = call_args[0][0]
        assert "<=>" in str(sql_obj)

    async def test_search_pgvector_with_filters(self):
        """pgvector search 应用过滤条件"""
        embedder = MockEmbedder(dimension=32)
        vec = await embedder.embed("test query")

        now = datetime.now(timezone.utc)

        mock_row = _make_row_mapping({
            "id": str(uuid.uuid4()),
            "agent_name": "specific_agent",
            "task_type": "analysis",
            "input_summary": "filtered result",
            "output_summary": "output",
            "outcome": "success",
            "quality_score": 0.8,
            "reflection": "",
            "embedding": vec,
            "created_at": now,
            "distance": 0.1,
        })

        mock_result = MagicMock()
        mock_result.mappings.return_value.all.return_value = [mock_row]

        mock_session = AsyncMock()
        mock_session.execute = AsyncMock(return_value=mock_result)

        @asynccontextmanager
        async def factory():
            yield mock_session

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

        results = await mem.search("test query", filters={"agent_name": "specific_agent"})
        assert len(results) == 1

        # Verify the SQL query contains WHERE clause
        call_args = mock_session.execute.call_args
        sql_obj = call_args[0][0]
        sql_text = str(sql_obj)
        assert "WHERE" in sql_text
        assert "agent_name" in sql_text

    async def test_search_pgvector_empty_result(self):
        """pgvector search 无结果时返回空列表"""
        embedder = MockEmbedder(dimension=32)

        mock_result = MagicMock()
        mock_result.mappings.return_value.all.return_value = []

        mock_session = AsyncMock()
        mock_session.execute = AsyncMock(return_value=mock_result)

        @asynccontextmanager
        async def factory():
            yield mock_session

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

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

    async def test_retrieve_pgvector_no_embedding_in_row(self):
        """pgvector retrieve 返回行没有 embedding 时返回 None"""
        embedder = MockEmbedder(dimension=32)

        mock_row = _make_row_mapping({
            "id": str(uuid.uuid4()),
            "embedding": None,
        })

        mock_result = MagicMock()
        mock_result.mappings.return_value.first.return_value = mock_row

        mock_session = AsyncMock()
        mock_session.execute = AsyncMock(return_value=mock_result)

        @asynccontextmanager
        async def factory():
            yield mock_session

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

        result = await mem.retrieve("test query")
        assert result is None

    async def test_retrieve_pgvector_no_rows(self):
        """pgvector retrieve 无匹配行时返回 None"""
        embedder = MockEmbedder(dimension=32)

        mock_result = MagicMock()
        mock_result.mappings.return_value.first.return_value = None

        mock_session = AsyncMock()
        mock_session.execute = AsyncMock(return_value=mock_result)

        @asynccontextmanager
        async def factory():
            yield mock_session

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

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

    async def test_search_pgvector_time_decay_reranking(self):
        """pgvector search 对返回结果做 time_decay 重排"""
        embedder = MockEmbedder(dimension=32)
        vec_similar = await embedder.embed("test query")
        vec_different = await embedder.embed("unrelated")

        now = datetime.now(timezone.utc)

        # Row with high cosine but low quality
        row_high_cosine = _make_row_mapping({
            "id": str(uuid.uuid4()),
            "agent_name": "",
            "task_type": "",
            "input_summary": "similar but low quality",
            "output_summary": "",
            "outcome": "success",
            "quality_score": 0.3,
            "reflection": "",
            "embedding": vec_similar,
            "created_at": now,
            "distance": 0.1,
        })

        # Row with lower cosine but high quality
        row_low_cosine = _make_row_mapping({
            "id": str(uuid.uuid4()),
            "agent_name": "",
            "task_type": "",
            "input_summary": "different but high quality",
            "output_summary": "",
            "outcome": "success",
            "quality_score": 0.9,
            "reflection": "",
            "embedding": vec_different,
            "created_at": now,
            "distance": 0.5,
        })

        mock_result = MagicMock()
        mock_result.mappings.return_value.all.return_value = [
            row_high_cosine,
            row_low_cosine,
        ]

        mock_session = AsyncMock()
        mock_session.execute = AsyncMock(return_value=mock_result)

        @asynccontextmanager
        async def factory():
            yield mock_session

        # alpha=1.0: pure cosine → similar entry first
        mem = EpisodicMemory(
            session_factory=factory,
            episodic_model=MockEpisodicModel,
            embedder=embedder,
            alpha=1.0,
            pgvector_enabled=True,
        )

        results = await mem.search("test query")
        assert len(results) == 2
        # With alpha=1.0, cosine dominates, so similar entry should be first
        assert results[0].value["input_summary"] == "similar but low quality"
