"""MultiSourceRAG 单元测试 - 多源混合检索

测试场景：
- 指定单个信息源 → 仅从该源检索
- 指定多个信息源 → 并行检索，结果融合排序
- 不指定信息源 → 从所有可用源检索
- 来源追溯 → 每个结果包含来源信息
- AE4: 指定"合规文档库"和"法务知识库" → 仅从这两个源检索
"""

import pytest

from agentkit.memory.embedder import MockEmbedder
from agentkit.memory.knowledge_base import Document, KnowledgeBase, QueryResult, SourceInfo
from agentkit.memory.local_rag import InMemoryLocalRAGService
from agentkit.memory.multi_source_retriever import MultiSourceRetriever
from agentkit.memory.retriever import MemoryRetriever


# ── Fixtures ──────────────────────────────────────────────


@pytest.fixture
def embedder():
    return MockEmbedder(dimension=128)


@pytest.fixture
def local_rag(embedder):
    """本地合规文档库"""
    return InMemoryLocalRAGService(embedder=embedder, chunk_size=500, chunk_overlap=50)


@pytest.fixture
def legal_rag(embedder):
    """法务知识库"""
    return InMemoryLocalRAGService(embedder=embedder, chunk_size=500, chunk_overlap=50)


@pytest.fixture
def tech_rag(embedder):
    """技术文档库"""
    return InMemoryLocalRAGService(embedder=embedder, chunk_size=500, chunk_overlap=50)


@pytest.fixture
def compliance_docs():
    """合规文档"""
    return [
        Document(
            doc_id="compliance-1",
            content="数据保护合规要求：所有用户数据必须加密存储，访问需经授权审批。",
            title="数据保护合规指南",
            source_id="compliance_data_protection",
            metadata={"source": "compliance_data_protection", "format": "text"},
        ),
        Document(
            doc_id="compliance-2",
            content="跨境数据传输需遵守 GDPR 和中国网络安全法的相关规定。",
            title="跨境数据传输合规",
            source_id="compliance_cross_border",
            metadata={"source": "compliance_cross_border", "format": "text"},
        ),
    ]


@pytest.fixture
def legal_docs():
    """法务文档"""
    return [
        Document(
            doc_id="legal-1",
            content="合同审查要点：注意违约责任条款、知识产权归属和保密义务。",
            title="合同审查指南",
            source_id="legal_contract_review",
            metadata={"source": "legal_contract_review", "format": "text"},
        ),
        Document(
            doc_id="legal-2",
            content="劳动法规定：员工加班需支付加班费，标准为平时工资的1.5倍至3倍。",
            title="劳动法要点",
            source_id="legal_labor_law",
            metadata={"source": "legal_labor_law", "format": "text"},
        ),
    ]


@pytest.fixture
def tech_docs():
    """技术文档"""
    return [
        Document(
            doc_id="tech-1",
            content="API 网关配置：限流策略为每分钟 1000 次请求，超时设置 30 秒。",
            title="API 网关配置手册",
            source_id="tech_api_gateway",
            metadata={"source": "tech_api_gateway", "format": "text"},
        ),
    ]


# ── MultiSourceRetriever 核心测试 ─────────────────────────


class TestMultiSourceRetrieverBasic:
    """MultiSourceRetriever 基础功能测试"""

    def test_register_source(self, local_rag, legal_rag):
        retriever = MultiSourceRetriever()
        retriever.register_source("local:合规文档", local_rag)
        retriever.register_source("法务知识库", legal_rag)

        names = retriever.get_source_names()
        assert "local:合规文档" in names
        assert "法务知识库" in names

    def test_register_source_via_constructor(self, local_rag, legal_rag):
        retriever = MultiSourceRetriever(
            sources={"local:合规文档": local_rag, "法务知识库": legal_rag}
        )

        names = retriever.get_source_names()
        assert len(names) == 2

    def test_unregister_source(self, local_rag, legal_rag):
        retriever = MultiSourceRetriever(
            sources={"local:合规文档": local_rag, "法务知识库": legal_rag}
        )

        result = retriever.unregister_source("local:合规文档")
        assert result is True
        assert "local:合规文档" not in retriever.get_source_names()

    def test_unregister_nonexistent_source(self, local_rag):
        retriever = MultiSourceRetriever(sources={"local:合规文档": local_rag})

        result = retriever.unregister_source("不存在")
        assert result is False

    @pytest.mark.asyncio
    async def test_list_all_sources(self, local_rag, legal_rag, compliance_docs, legal_docs):
        await local_rag.ingest(compliance_docs)
        await legal_rag.ingest(legal_docs)

        retriever = MultiSourceRetriever(
            sources={"local:合规文档": local_rag, "法务知识库": legal_rag}
        )

        sources = await retriever.list_all_sources()
        assert "local:合规文档" in sources
        assert "法务知识库" in sources
        for name, info in sources.items():
            assert isinstance(info, SourceInfo)


class TestMultiSourceRetrieverSearch:
    """MultiSourceRetriever 检索功能测试"""

    @pytest.mark.asyncio
    async def test_search_single_source(
        self, local_rag, legal_rag, compliance_docs, legal_docs
    ):
        """指定单个信息源 → 仅从该源检索"""
        await local_rag.ingest(compliance_docs)
        await legal_rag.ingest(legal_docs)

        retriever = MultiSourceRetriever(
            sources={"local:合规文档": local_rag, "法务知识库": legal_rag}
        )

        # 仅从合规文档库检索
        results = await retriever.search("合规", top_k=5, sources=["local:合规文档"])

        # 所有结果应来自合规文档库
        for r in results:
            assert r.source_name == "local:合规文档"

    @pytest.mark.asyncio
    async def test_search_multiple_sources(
        self, local_rag, legal_rag, compliance_docs, legal_docs
    ):
        """指定多个信息源 → 并行检索，结果融合排序"""
        await local_rag.ingest(compliance_docs)
        await legal_rag.ingest(legal_docs)

        retriever = MultiSourceRetriever(
            sources={"local:合规文档": local_rag, "法务知识库": legal_rag}
        )

        results = await retriever.search(
            "合规 法务", top_k=10, sources=["local:合规文档", "法务知识库"]
        )

        # 结果应来自两个源
        source_names = {r.source_name for r in results}
        assert source_names.issubset({"local:合规文档", "法务知识库"})

    @pytest.mark.asyncio
    async def test_search_all_sources_when_none_specified(
        self, local_rag, legal_rag, tech_rag, compliance_docs, legal_docs, tech_docs
    ):
        """不指定信息源 → 从所有可用源检索"""
        await local_rag.ingest(compliance_docs)
        await legal_rag.ingest(legal_docs)
        await tech_rag.ingest(tech_docs)

        retriever = MultiSourceRetriever(
            sources={
                "local:合规文档": local_rag,
                "法务知识库": legal_rag,
                "技术文档库": tech_rag,
            }
        )

        results = await retriever.search("合规 法务 技术", top_k=10)

        # 结果应来自所有三个源
        source_names = {r.source_name for r in results}
        assert len(source_names) >= 1  # 至少有一个源返回结果

    @pytest.mark.asyncio
    async def test_search_no_sources_registered(self):
        """无信息源注册时返回空结果"""
        retriever = MultiSourceRetriever()

        results = await retriever.search("anything", top_k=5)
        assert len(results) == 0

    @pytest.mark.asyncio
    async def test_search_nonexistent_source(self, local_rag, compliance_docs):
        """指定不存在的源 → 跳过，返回空结果"""
        await local_rag.ingest(compliance_docs)

        retriever = MultiSourceRetriever(sources={"local:合规文档": local_rag})

        results = await retriever.search("合规", top_k=5, sources=["不存在的源"])
        assert len(results) == 0

    @pytest.mark.asyncio
    async def test_search_with_weights(
        self, local_rag, legal_rag, compliance_docs, legal_docs
    ):
        """带权重检索 → 特定源分数被调整"""
        await local_rag.ingest(compliance_docs)
        await legal_rag.ingest(legal_docs)

        retriever = MultiSourceRetriever(
            sources={"local:合规文档": local_rag, "法务知识库": legal_rag}
        )

        # 先不带权重检索
        results_no_weight = await retriever.search(
            "合规", top_k=10, sources=["local:合规文档", "法务知识库"]
        )

        # 带权重检索：提升合规文档库
        results_with_weight = await retriever.search(
            "合规",
            top_k=10,
            sources=["local:合规文档", "法务知识库"],
            weights={"local:合规文档": 2.0},
        )

        # 有权重时合规文档库的分数应更高
        compliance_scores_weighted = [
            r.score for r in results_with_weight if r.source_name == "local:合规文档"
        ]
        compliance_scores_unweighted = [
            r.score for r in results_no_weight if r.source_name == "local:合规文档"
        ]

        if compliance_scores_weighted and compliance_scores_unweighted:
            assert max(compliance_scores_weighted) >= max(compliance_scores_unweighted)


class TestMultiSourceRetrieverSourceTracing:
    """来源追溯测试"""

    @pytest.mark.asyncio
    async def test_result_contains_source_info(
        self, local_rag, legal_rag, compliance_docs, legal_docs
    ):
        """每个检索结果包含来源追溯信息"""
        await local_rag.ingest(compliance_docs)
        await legal_rag.ingest(legal_docs)

        retriever = MultiSourceRetriever(
            sources={"local:合规文档": local_rag, "法务知识库": legal_rag}
        )

        results = await retriever.search("合规", top_k=5)

        for r in results:
            # source_id 应非空
            assert r.source_id != ""
            # source_name 应为注册的源名称
            assert r.source_name in ("local:合规文档", "法务知识库")
            # title 应非空
            assert r.title != ""

    @pytest.mark.asyncio
    async def test_result_contains_document_title(
        self, local_rag, compliance_docs
    ):
        """检索结果包含文档标题"""
        await local_rag.ingest(compliance_docs)

        retriever = MultiSourceRetriever(sources={"local:合规文档": local_rag})

        results = await retriever.search("数据保护", top_k=5)

        for r in results:
            assert r.title != ""
            assert r.doc_id != ""


class TestMultiSourceRetrieverDedup:
    """去重测试"""

    @pytest.mark.asyncio
    async def test_deduplicate_identical_content(
        self, local_rag, legal_rag, embedder
    ):
        """相同内容从不同源返回时去重，保留高分"""
        # 两个源包含相同内容
        same_doc = Document(
            doc_id="same-doc",
            content="这是一段完全相同的内容用于测试去重功能。",
            title="重复文档",
            source_id="same_source",
            metadata={"source": "same_source", "format": "text"},
        )
        await local_rag.ingest([same_doc])
        await legal_rag.ingest([same_doc])

        retriever = MultiSourceRetriever(
            sources={"local:合规文档": local_rag, "法务知识库": legal_rag}
        )

        results = await retriever.search("去重", top_k=10)

        # 相同内容应去重，只保留一个
        content_counts: dict[str, int] = {}
        for r in results:
            content_counts[r.content] = content_counts.get(r.content, 0) + 1

        for content, count in content_counts.items():
            assert count == 1, f"内容 '{content[:30]}...' 出现了 {count} 次，应去重为 1 次"


# ── AE4: 合规文档库 + 法务知识库指定检索 ──────────────────


class TestAE4ComplianceAndLegalSearch:
    """AE4 场景：指定"合规文档库"和"法务知识库" → 仅从这两个源检索"""

    @pytest.mark.asyncio
    async def test_search_compliance_and_legal_only(
        self, local_rag, legal_rag, tech_rag, compliance_docs, legal_docs, tech_docs
    ):
        """指定合规和法务源 → 不从技术文档库检索"""
        await local_rag.ingest(compliance_docs)
        await legal_rag.ingest(legal_docs)
        await tech_rag.ingest(tech_docs)

        retriever = MultiSourceRetriever(
            sources={
                "合规文档库": local_rag,
                "法务知识库": legal_rag,
                "技术文档库": tech_rag,
            }
        )

        results = await retriever.search(
            "合规 法务", top_k=10, sources=["合规文档库", "法务知识库"]
        )

        # 结果不应来自技术文档库
        for r in results:
            assert r.source_name != "技术文档库"
            assert r.source_name in ("合规文档库", "法务知识库")

    @pytest.mark.asyncio
    async def test_search_compliance_and_legal_results_merged(
        self, local_rag, legal_rag, compliance_docs, legal_docs
    ):
        """合规和法务源的结果应合并排序"""
        await local_rag.ingest(compliance_docs)
        await legal_rag.ingest(legal_docs)

        retriever = MultiSourceRetriever(
            sources={"合规文档库": local_rag, "法务知识库": legal_rag}
        )

        results = await retriever.search(
            "合规 法务", top_k=10, sources=["合规文档库", "法务知识库"]
        )

        # 应有来自两个源的结果
        source_names = {r.source_name for r in results}
        assert len(source_names) >= 1

        # 结果应按 score 降序排列
        for i in range(len(results) - 1):
            assert results[i].score >= results[i + 1].score


# ── MemoryRetriever 集成测试 ──────────────────────────────


class TestMemoryRetrieverIntegration:
    """MemoryRetriever 与 MultiSourceRetriever 集成测试"""

    @pytest.mark.asyncio
    async def test_retrieve_with_sources_parameter(
        self, local_rag, legal_rag, compliance_docs, legal_docs
    ):
        """MemoryRetriever.retrieve(sources=...) 委托给 MultiSourceRetriever"""
        await local_rag.ingest(compliance_docs)
        await legal_rag.ingest(legal_docs)

        retriever = MemoryRetriever(
            knowledge_sources={"local:合规文档": local_rag, "法务知识库": legal_rag}
        )

        items = await retriever.retrieve(
            "合规", top_k=5, sources=["local:合规文档"]
        )

        # 结果应为 MemoryItem 类型
        for item in items:
            assert hasattr(item, "key")
            assert hasattr(item, "value")
            assert hasattr(item, "score")
            assert hasattr(item, "metadata")

    @pytest.mark.asyncio
    async def test_retrieve_without_sources_keeps_current_behavior(
        self, local_rag, compliance_docs
    ):
        """不指定 sources 时保持原有行为（三层记忆检索）"""
        await local_rag.ingest(compliance_docs)

        retriever = MemoryRetriever(
            knowledge_sources={"local:合规文档": local_rag}
        )

        # 不指定 sources → 走三层记忆路径
        items = await retriever.retrieve("合规", top_k=5)
        # 三层记忆为空，应返回空结果
        assert isinstance(items, list)

    @pytest.mark.asyncio
    async def test_retrieve_from_sources_with_source_tracing(
        self, local_rag, legal_rag, compliance_docs, legal_docs
    ):
        """通过 MemoryRetriever 多源检索时，结果包含来源追溯"""
        await local_rag.ingest(compliance_docs)
        await legal_rag.ingest(legal_docs)

        retriever = MemoryRetriever(
            knowledge_sources={"local:合规文档": local_rag, "法务知识库": legal_rag}
        )

        items = await retriever.retrieve(
            "合规", top_k=5, sources=["local:合规文档", "法务知识库"]
        )

        for item in items:
            assert item.metadata.get("source") == "rag"
            assert "source_name" in item.metadata
            assert "document_title" in item.metadata

    @pytest.mark.asyncio
    async def test_multi_source_retriever_property(
        self, local_rag, compliance_docs
    ):
        """通过 multi_source_retriever 属性直接访问"""
        retriever = MemoryRetriever(
            knowledge_sources={"local:合规文档": local_rag}
        )

        ms_retriever = retriever.multi_source_retriever
        assert isinstance(ms_retriever, MultiSourceRetriever)
        assert "local:合规文档" in ms_retriever.get_source_names()

    @pytest.mark.asyncio
    async def test_register_source_via_property(
        self, local_rag, legal_rag, compliance_docs, legal_docs
    ):
        """通过 multi_source_retriever 属性动态注册源"""
        await local_rag.ingest(compliance_docs)
        await legal_rag.ingest(legal_docs)

        retriever = MemoryRetriever(
            knowledge_sources={"local:合规文档": local_rag}
        )

        # 动态注册法务知识库
        retriever.multi_source_retriever.register_source("法务知识库", legal_rag)

        # 现在可以从法务知识库检索
        items = await retriever.retrieve(
            "合同", top_k=5, sources=["法务知识库"]
        )

        for item in items:
            assert item.metadata.get("source_name") == "法务知识库"

    @pytest.mark.asyncio
    async def test_retrieve_with_source_weights(
        self, local_rag, legal_rag, compliance_docs, legal_docs
    ):
        """MemoryRetriever 支持 source_weights 参数"""
        await local_rag.ingest(compliance_docs)
        await legal_rag.ingest(legal_docs)

        retriever = MemoryRetriever(
            knowledge_sources={"local:合规文档": local_rag, "法务知识库": legal_rag}
        )

        items = await retriever.retrieve(
            "合规",
            top_k=5,
            sources=["local:合规文档", "法务知识库"],
            source_weights={"local:合规文档": 1.5},
        )

        assert isinstance(items, list)


# ── 边界情况测试 ──────────────────────────────────────────


class TestEdgeCases:
    """边界情况测试"""

    @pytest.mark.asyncio
    async def test_search_with_empty_source_list(self, local_rag, compliance_docs):
        """sources=[] 空列表 → 不查询任何源，返回空"""
        await local_rag.ingest(compliance_docs)

        retriever = MultiSourceRetriever(sources={"local:合规文档": local_rag})

        results = await retriever.search("合规", top_k=5, sources=[])
        assert len(results) == 0

    @pytest.mark.asyncio
    async def test_search_top_k_limits_results(
        self, local_rag, legal_rag, compliance_docs, legal_docs
    ):
        """top_k 限制返回结果数量"""
        await local_rag.ingest(compliance_docs)
        await legal_rag.ingest(legal_docs)

        retriever = MultiSourceRetriever(
            sources={"local:合规文档": local_rag, "法务知识库": legal_rag}
        )

        results = await retriever.search("合规", top_k=1)
        assert len(results) <= 1

    @pytest.mark.asyncio
    async def test_source_query_failure_graceful(
        self, local_rag, compliance_docs
    ):
        """某个源查询失败时，其他源结果正常返回"""
        await local_rag.ingest(compliance_docs)

        # 创建一个会抛异常的 mock 源
        class FailingSource:
            async def ingest(self, documents):
                return []

            async def query(self, text, top_k=5):
                raise ConnectionError("Service unavailable")

            async def delete_by_id(self, id):
                return False

            async def list_sources(self):
                return [SourceInfo(source_id="failing", source_name="Failing", source_type="mock")]

            async def health_check(self):
                return False

        retriever = MultiSourceRetriever(
            sources={"local:合规文档": local_rag, "failing_source": FailingSource()}
        )

        # 应不抛异常，返回合规文档库的结果
        results = await retriever.search("合规", top_k=5)
        assert isinstance(results, list)

    @pytest.mark.asyncio
    async def test_search_results_sorted_by_score(
        self, local_rag, legal_rag, compliance_docs, legal_docs
    ):
        """结果按 score 降序排列"""
        await local_rag.ingest(compliance_docs)
        await legal_rag.ingest(legal_docs)

        retriever = MultiSourceRetriever(
            sources={"local:合规文档": local_rag, "法务知识库": legal_rag}
        )

        results = await retriever.search("合规 法务", top_k=10)

        for i in range(len(results) - 1):
            assert results[i].score >= results[i + 1].score