"""U4 测试 — 双索引检索引擎（embedding/keywords/blend）。

测试场景：
1. embedding 模式：语义检索返回相关结果，按 kb_id 过滤
2. keywords 模式：中文全文检索返回包含关键词的结果
3. blend 模式：合并语义+全文结果，去重排序
4. 查询无结果时返回空列表（非报错）
5. RetrievalRequest 模型字段
"""

from __future__ import annotations

import sys
from contextlib import asynccontextmanager
from unittest.mock import AsyncMock, MagicMock

from agentkit.rag_platform.models import QueryMode, QueryResult, RetrievalRequest
from agentkit.rag_platform.retrieval import RetrievalEngine, _normalize_scores


# ---------------------------------------------------------------------------
# llama_index 模块 mock — 测试环境可能未安装 llama_index
# ---------------------------------------------------------------------------


def _setup_llama_index_mocks():
    """注入 mock llama_index 模块到 sys.modules，使 import 成功。

    仅当真实 llama_index 无法导入时才注入 mock — 避免污染已安装真实模块的环境。
    """
    try:
        import llama_index  # noqa: F401

        return  # 真实模块可用，无需 mock
    except ImportError:
        pass

    # 创建 mock 模块层级
    mock_li = MagicMock()
    mock_li_core = MagicMock()
    mock_li_vector_stores = MagicMock()
    mock_li_core_vector_stores_types = MagicMock()
    mock_li_core_embeddings = MagicMock()

    # VectorStoreQuery — 简单的可调用 mock
    mock_li_core_vector_stores_types.VectorStoreQuery = MagicMock()

    sys.modules["llama_index"] = mock_li
    sys.modules["llama_index.core"] = mock_li_core
    sys.modules["llama_index.vector_stores"] = mock_li_vector_stores
    sys.modules["llama_index.core.vector_stores"] = mock_li_core
    sys.modules["llama_index.core.vector_stores.types"] = mock_li_core_vector_stores_types
    sys.modules["llama_index.core.embeddings"] = mock_li_core_embeddings


_setup_llama_index_mocks()


# ---------------------------------------------------------------------------
# 测试辅助函数
# ---------------------------------------------------------------------------


def _make_mock_embed_model():
    """创建 mock embedding 模型。"""
    mock = MagicMock()
    mock.aget_text_embedding = AsyncMock(return_value=[0.1] * 1536)
    return mock


def _make_mock_text_node(node_id: str, content: str, metadata: dict | None = None):
    """创建 mock LlamaIndex TextNode。"""
    node = MagicMock()
    node.node_id = node_id
    node.get_content.return_value = content
    node.metadata = metadata or {}
    return node


def _make_mock_vector_store(nodes=None, similarities=None):
    """创建 mock PGVectorStore。"""
    mock = MagicMock()
    mock.aquery = AsyncMock()
    if nodes is not None:
        mock_result = MagicMock()
        mock_result.nodes = nodes
        mock_result.similarities = similarities or [0.9] * len(nodes)
        mock.aquery.return_value = mock_result
    return mock


def _make_mock_session_factory(execute_result_rows=None):
    """创建 mock session factory，用于 keywords 模式测试。

    Args:
        execute_result_rows: list of tuples — SQL 查询返回的行
            每行格式: (id, text, metadata_, document_id, score)
    """
    mock_session = AsyncMock()

    mock_result = MagicMock()
    mock_result.all.return_value = execute_result_rows or []
    mock_session.execute = AsyncMock(return_value=mock_result)

    @asynccontextmanager
    async def factory():
        yield mock_session

    return factory, mock_session


# ---------------------------------------------------------------------------
# RetrievalRequest 模型测试
# ---------------------------------------------------------------------------


class TestRetrievalRequest:
    """RetrievalRequest 模型测试。"""

    def test_defaults(self):
        """默认值正确。"""
        req = RetrievalRequest(query="test", kb_ids=["kb1"])
        assert req.query == "test"
        assert req.kb_ids == ["kb1"]
        assert req.retrieval_mode is None
        assert req.hit_processing_mode is None
        assert req.top_k == 5
        assert req.user_id is None

    def test_explicit_values(self):
        """显式赋值正确。"""
        req = RetrievalRequest(
            query="hello",
            kb_ids=["kb1", "kb2"],
            retrieval_mode=QueryMode.embedding,
            hit_processing_mode="direct",
            top_k=10,
            user_id="user1",
        )
        assert req.retrieval_mode == QueryMode.embedding
        assert req.hit_processing_mode == "direct"
        assert req.top_k == 10
        assert req.user_id == "user1"


# ---------------------------------------------------------------------------
# _normalize_scores 测试
# ---------------------------------------------------------------------------


class TestNormalizeScores:
    """分数归一化测试。"""

    def test_empty_list(self):
        """空列表返回空列表。"""
        assert _normalize_scores([]) == []

    def test_constant_list_returns_ones(self):
        """常数列表（含单元素）返回全 1.0 — 等价于最高相关度。"""
        assert _normalize_scores([0.5, 0.5, 0.5]) == [1.0, 1.0, 1.0]
        assert _normalize_scores([0.9]) == [1.0]

    def test_min_max_normalization(self):
        """min-max 归一化到 [0, 1]。"""
        result = _normalize_scores([1.0, 2.0, 3.0])
        assert result[0] == 0.0  # 最小值
        assert result[2] == 1.0  # 最大值
        assert 0.0 < result[1] < 1.0


# ---------------------------------------------------------------------------
# embedding 模式测试
# ---------------------------------------------------------------------------


class TestEmbeddingRetrieval:
    """embedding 模式检索测试。"""

    async def test_returns_relevant_results(self):
        """embedding 模式返回相关结果。"""
        nodes = [
            _make_mock_text_node("n1", "chunk 1", {"kb_id": "kb1", "document_id": "d1"}),
            _make_mock_text_node("n2", "chunk 2", {"kb_id": "kb1", "document_id": "d1"}),
        ]
        mock_vs = _make_mock_vector_store(nodes, [0.95, 0.80])
        mock_embed = _make_mock_embed_model()
        sf, _ = _make_mock_session_factory()

        engine = RetrievalEngine(mock_vs, sf, mock_embed)
        results = await engine.retrieve("query", ["kb1"], QueryMode.embedding, top_k=5)

        assert len(results) == 2
        assert all(isinstance(r, QueryResult) for r in results)
        assert results[0].chunk_id == "n1"
        assert results[0].score == 0.95
        assert results[0].kb_id == "kb1"

    async def test_filters_by_kb_id(self):
        """embedding 模式按 kb_id 过滤结果。"""
        # n1 属于 kb1，n2 属于 kb2（不在查询范围）
        nodes = [
            _make_mock_text_node("n1", "chunk 1", {"kb_id": "kb1"}),
            _make_mock_text_node("n2", "chunk 2", {"kb_id": "kb2"}),
        ]
        mock_vs = _make_mock_vector_store(nodes, [0.9, 0.9])
        mock_embed = _make_mock_embed_model()
        sf, _ = _make_mock_session_factory()

        engine = RetrievalEngine(mock_vs, sf, mock_embed)
        results = await engine.retrieve("query", ["kb1"], QueryMode.embedding, top_k=5)

        assert len(results) == 1
        assert results[0].kb_id == "kb1"

    async def test_empty_results_no_error(self):
        """无结果时返回空列表（非报错）。"""
        mock_vs = _make_mock_vector_store([], [])
        mock_embed = _make_mock_embed_model()
        sf, _ = _make_mock_session_factory()

        engine = RetrievalEngine(mock_vs, sf, mock_embed)
        results = await engine.retrieve("query", ["kb1"], QueryMode.embedding, top_k=5)

        assert results == []

    async def test_empty_kb_ids_returns_empty(self):
        """kb_ids 为空时直接返回空列表。"""
        mock_vs = _make_mock_vector_store()
        mock_embed = _make_mock_embed_model()
        sf, _ = _make_mock_session_factory()

        engine = RetrievalEngine(mock_vs, sf, mock_embed)
        results = await engine.retrieve("query", [], QueryMode.embedding, top_k=5)

        assert results == []
        mock_vs.aquery.assert_not_awaited()

    async def test_no_embed_model_raises(self):
        """embedding 模式未配置 embed_model 抛出异常。"""
        mock_vs = _make_mock_vector_store()
        sf, _ = _make_mock_session_factory()

        engine = RetrievalEngine(mock_vs, sf, embed_model=None)
        try:
            await engine.retrieve("query", ["kb1"], QueryMode.embedding, top_k=5)
            raise AssertionError("Expected ValueError")
        except ValueError as e:
            assert "embed_model" in str(e)


# ---------------------------------------------------------------------------
# keywords 模式测试
# ---------------------------------------------------------------------------


class TestKeywordsRetrieval:
    """keywords 模式检索测试。"""

    async def test_returns_matching_results(self):
        """keywords 模式返回包含关键词的结果。"""
        rows = [
            ("c1", "自然语言处理内容", {"kb_id": "kb1", "document_id": "d1"}, "d1", 0.8),
            ("c2", "另一段文本", {"kb_id": "kb1", "document_id": "d1"}, "d1", 0.5),
        ]
        sf, _ = _make_mock_session_factory(rows)
        mock_vs = _make_mock_vector_store()

        engine = RetrievalEngine(mock_vs, sf, embed_model=None)
        results = await engine.retrieve("自然语言", ["kb1"], QueryMode.keywords, top_k=5)

        assert len(results) == 2
        assert results[0].chunk_id == "c1"
        assert results[0].score == 0.8
        assert results[0].kb_id == "kb1"
        assert "自然语言" in results[0].content

    async def test_empty_results_no_error(self):
        """无结果时返回空列表（非报错）。"""
        sf, _ = _make_mock_session_factory([])
        mock_vs = _make_mock_vector_store()

        engine = RetrievalEngine(mock_vs, sf, embed_model=None)
        results = await engine.retrieve("query", ["kb1"], QueryMode.keywords, top_k=5)

        assert results == []

    async def test_empty_query_returns_empty(self):
        """空查询（分词后无 token）返回空列表，不执行 SQL。"""
        sf, mock_session = _make_mock_session_factory([])
        mock_vs = _make_mock_vector_store()

        engine = RetrievalEngine(mock_vs, sf, embed_model=None)
        results = await engine.retrieve("   ", ["kb1"], QueryMode.keywords, top_k=5)

        assert results == []
        mock_session.execute.assert_not_awaited()

    async def test_passes_kb_ids_to_sql(self):
        """kb_ids 作为参数传递给 SQL 查询。"""
        rows = [
            ("c1", "content", {"kb_id": "kb1"}, "d1", 0.5),
        ]
        sf, mock_session = _make_mock_session_factory(rows)
        mock_vs = _make_mock_vector_store()

        engine = RetrievalEngine(mock_vs, sf, embed_model=None)
        await engine.retrieve("test", ["kb1", "kb2"], QueryMode.keywords, top_k=3)

        call_args = mock_session.execute.await_args
        params = call_args.args[1] if len(call_args.args) > 1 else call_args.kwargs
        assert params["kb_ids"] == ["kb1", "kb2"]
        assert params["top_k"] == 3

    async def test_handles_none_metadata(self):
        """metadata 为 None 时正常处理。"""
        rows = [
            ("c1", "content", None, "d1", 0.5),
        ]
        sf, _ = _make_mock_session_factory(rows)
        mock_vs = _make_mock_vector_store()

        engine = RetrievalEngine(mock_vs, sf, embed_model=None)
        results = await engine.retrieve("test", ["kb1"], QueryMode.keywords, top_k=5)

        assert len(results) == 1
        assert results[0].metadata == {}
        assert results[0].kb_id == ""


# ---------------------------------------------------------------------------
# blend 模式测试
# ---------------------------------------------------------------------------


class TestBlendRetrieval:
    """blend 模式检索测试。"""

    async def test_merges_and_deduplicates(self):
        """blend 模式合并语义+全文结果，按 chunk_id 去重。"""
        # embedding 结果
        nodes = [
            _make_mock_text_node("c1", "chunk 1", {"kb_id": "kb1", "document_id": "d1"}),
            _make_mock_text_node("c2", "chunk 2", {"kb_id": "kb1", "document_id": "d1"}),
        ]
        mock_vs = _make_mock_vector_store(nodes, [0.9, 0.7])
        mock_embed = _make_mock_embed_model()

        # keywords 结果 — c1 重复，c3 新增
        rows = [
            ("c1", "chunk 1", {"kb_id": "kb1", "document_id": "d1"}, "d1", 0.8),
            ("c3", "chunk 3", {"kb_id": "kb1", "document_id": "d1"}, "d1", 0.6),
        ]
        sf, _ = _make_mock_session_factory(rows)

        engine = RetrievalEngine(mock_vs, sf, mock_embed)
        results = await engine.retrieve("query", ["kb1"], QueryMode.blend, top_k=5)

        # 去重后应有 3 个结果（c1, c2, c3）
        chunk_ids = {r.chunk_id for r in results}
        assert chunk_ids == {"c1", "c2", "c3"}
        # 结果按分数降序
        scores = [r.score for r in results]
        assert scores == sorted(scores, reverse=True)

    async def test_dedup_takes_highest_score(self):
        """去重时同 chunk_id 取最高分。"""
        # embedding: c1 score=0.9 (归一化后 1.0 * 0.6 = 0.6)
        nodes = [
            _make_mock_text_node("c1", "chunk 1", {"kb_id": "kb1"}),
        ]
        mock_vs = _make_mock_vector_store(nodes, [0.9])
        mock_embed = _make_mock_embed_model()

        # keywords: c1 score=0.8 (归一化后 1.0 * 0.4 = 0.4)
        rows = [
            ("c1", "chunk 1", {"kb_id": "kb1"}, "d1", 0.8),
        ]
        sf, _ = _make_mock_session_factory(rows)

        engine = RetrievalEngine(mock_vs, sf, mock_embed)
        results = await engine.retrieve("query", ["kb1"], QueryMode.blend, top_k=5)

        # c1 只出现一次，取 embedding 的分数（0.6 > 0.4）
        assert len(results) == 1
        assert results[0].chunk_id == "c1"
        assert results[0].score == 0.6  # 1.0 * 0.6

    async def test_empty_results_no_error(self):
        """两种检索都无结果时返回空列表。"""
        mock_vs = _make_mock_vector_store([], [])
        mock_embed = _make_mock_embed_model()
        sf, _ = _make_mock_session_factory([])

        engine = RetrievalEngine(mock_vs, sf, mock_embed)
        results = await engine.retrieve("query", ["kb1"], QueryMode.blend, top_k=5)

        assert results == []

    async def test_respects_top_k(self):
        """blend 模式遵守 top_k 限制。"""
        # embedding 返回 3 个结果
        nodes = [_make_mock_text_node(f"n{i}", f"chunk {i}", {"kb_id": "kb1"}) for i in range(3)]
        mock_vs = _make_mock_vector_store(nodes, [0.9, 0.8, 0.7])
        mock_embed = _make_mock_embed_model()
        sf, _ = _make_mock_session_factory([])  # keywords 无结果

        engine = RetrievalEngine(mock_vs, sf, mock_embed)
        results = await engine.retrieve("query", ["kb1"], QueryMode.blend, top_k=2)

        assert len(results) == 2

    async def test_fallback_when_embedding_fails(self):
        """embedding 检索失败时降级为纯关键词检索。"""
        # embedding 模式会抛异常（mock 设置）
        mock_vs = MagicMock()
        mock_vs.aquery = AsyncMock(side_effect=RuntimeError("connection failed"))
        mock_embed = _make_mock_embed_model()

        # keywords 返回结果
        rows = [
            ("c1", "content", {"kb_id": "kb1"}, "d1", 0.5),
        ]
        sf, _ = _make_mock_session_factory(rows)

        engine = RetrievalEngine(mock_vs, sf, mock_embed)
        results = await engine.retrieve("query", ["kb1"], QueryMode.blend, top_k=5)

        # 应降级为纯关键词结果
        assert len(results) == 1
        assert results[0].chunk_id == "c1"

    async def test_fallback_when_no_embed_model(self):
        """未配置 embed_model 时 blend 降级为纯关键词检索。"""
        mock_vs = _make_mock_vector_store()
        mock_embed = None

        rows = [
            ("c1", "content", {"kb_id": "kb1"}, "d1", 0.5),
        ]
        sf, _ = _make_mock_session_factory(rows)

        engine = RetrievalEngine(mock_vs, sf, mock_embed)
        results = await engine.retrieve("query", ["kb1"], QueryMode.blend, top_k=5)

        assert len(results) == 1
        assert results[0].chunk_id == "c1"


# ---------------------------------------------------------------------------
# 不支持的 mode 测试
# ---------------------------------------------------------------------------


class TestUnsupportedMode:
    """不支持的检索模式测试。"""

    async def test_unsupported_mode_raises(self):
        """不支持的 mode 抛出 ValueError。"""
        mock_vs = _make_mock_vector_store()
        sf, _ = _make_mock_session_factory()

        engine = RetrievalEngine(mock_vs, sf, embed_model=None)
        try:
            await engine.retrieve("query", ["kb1"], "invalid_mode", top_k=5)  # type: ignore[arg-type]
            raise AssertionError("Expected ValueError")
        except ValueError as e:
            assert "Unsupported" in str(e)
