"""MultiSourceRetriever - 多源混合检索器

管理多个 KnowledgeBase 协议实现，支持：
- 信息源指定：search(query, sources=["feishu", "local:合规文档"])
- 并行查询多个源，按权重融合排序
- 来源追溯：每个检索结果附带 source_id + document_title + chunk_location
- 基于 hash 的去重
"""

from __future__ import annotations

import asyncio
import hashlib
import logging
from dataclasses import replace
from typing import Any

from agentkit.memory.knowledge_base import KnowledgeBase, QueryResult, SourceInfo

logger = logging.getLogger(__name__)


def _content_hash(content: str) -> str:
    """计算内容哈希，用于去重"""
    return hashlib.md5(content.encode("utf-8")).hexdigest()


class MultiSourceRetriever:
    """多源混合检索器

    管理多个 KnowledgeBase 协议实现，支持按名称指定信息源、
    并行查询、权重融合排序和来源追溯。

    用法::

        retriever = MultiSourceRetriever()
        retriever.register_source("feishu", feishu_adapter)
        retriever.register_source("local:合规文档", local_rag)

        # 仅从指定源检索
        results = await retriever.search("合规要求", sources=["feishu", "local:合规文档"])

        # 从所有可用源检索
        results = await retriever.search("合规要求")

        # 带权重检索
        results = await retriever.search("合规要求", weights={"feishu": 1.5, "local:合规文档": 0.8})
    """

    def __init__(self, sources: dict[str, KnowledgeBase] | None = None):
        """
        Args:
            sources: 初始信息源映射，key 为源名称，value 为 KnowledgeBase 实现
        """
        self._sources: dict[str, KnowledgeBase] = {}
        if sources:
            for name, kb in sources.items():
                self._sources[name] = kb

    def register_source(self, name: str, knowledge_base: KnowledgeBase) -> None:
        """注册信息源

        Args:
            name: 信息源名称，如 "feishu"、"local:合规文档"
            knowledge_base: KnowledgeBase 协议实现
        """
        self._sources[name] = knowledge_base
        logger.info(f"Registered knowledge source: {name}")

    def unregister_source(self, name: str) -> bool:
        """注销信息源

        Args:
            name: 信息源名称

        Returns:
            是否成功注销
        """
        if name in self._sources:
            del self._sources[name]
            logger.info(f"Unregistered knowledge source: {name}")
            return True
        return False

    async def search(
        self,
        query: str,
        top_k: int = 5,
        sources: list[str] | None = None,
        weights: dict[str, float] | None = None,
    ) -> list[QueryResult]:
        """多源检索

        Args:
            query: 检索查询
            top_k: 返回最大结果数
            sources: 指定信息源列表，None 表示查询所有已注册源
            weights: 信息源权重映射，用于提升/降低特定源的分数

        Returns:
            融合排序后的检索结果列表，每个结果包含来源追溯信息
        """
        # 确定要查询的源
        target_sources = self._resolve_sources(sources)
        if not target_sources:
            logger.warning("No knowledge sources available for search")
            return []

        # 并行查询所有目标源
        results = await self._query_sources(query, top_k, target_sources, weights)

        # 去重
        results = self._deduplicate(results)

        # 按 score 降序排序，截取 top_k
        results.sort(key=lambda r: r.score, reverse=True)
        return results[:top_k]

    async def list_all_sources(self) -> dict[str, SourceInfo]:
        """列出所有已注册信息源

        Returns:
            源名称到 SourceInfo 的映射
        """
        result: dict[str, SourceInfo] = {}
        for name, kb in self._sources.items():
            try:
                source_infos = await kb.list_sources()
                if source_infos:
                    result[name] = source_infos[0]
                else:
                    # 知识库未返回 source info，构造一个占位
                    result[name] = SourceInfo(
                        source_id=name,
                        source_name=name,
                        source_type="unknown",
                    )
            except Exception as e:
                logger.error(f"Failed to list sources for '{name}': {e}")
                result[name] = SourceInfo(
                    source_id=name,
                    source_name=name,
                    source_type="error",
                )
        return result

    def get_source_names(self) -> list[str]:
        """获取所有已注册信息源名称"""
        return list(self._sources.keys())

    def _resolve_sources(self, sources: list[str] | None) -> dict[str, KnowledgeBase]:
        """解析目标信息源

        Args:
            sources: 指定的源名称列表，None 表示所有源

        Returns:
            源名称到 KnowledgeBase 的映射
        """
        if sources is None:
            return dict(self._sources)

        resolved = {}
        for name in sources:
            if name in self._sources:
                resolved[name] = self._sources[name]
            else:
                logger.warning(f"Knowledge source '{name}' not found, skipping")
        return resolved

    async def _query_sources(
        self,
        query: str,
        top_k: int,
        target_sources: dict[str, KnowledgeBase],
        weights: dict[str, float] | None,
    ) -> list[QueryResult]:
        """并行查询多个信息源

        Args:
            query: 查询文本
            top_k: 每个源返回的最大结果数
            target_sources: 目标源映射
            weights: 权重映射

        Returns:
            所有源的检索结果列表（已应用权重）
        """
        async def _query_one(name: str, kb: KnowledgeBase) -> list[QueryResult]:
            try:
                results = await kb.query(query, top_k=top_k)
                # 应用权重
                weight = (weights or {}).get(name, 1.0)
                return [
                    replace(r, score=r.score * weight, source_name=name)
                    for r in results
                ]
            except Exception as e:
                logger.error(f"Query failed for source '{name}': {e}")
                return []

        tasks = [_query_one(name, kb) for name, kb in target_sources.items()]
        all_results = await asyncio.gather(*tasks, return_exceptions=True)

        merged: list[QueryResult] = []
        for result in all_results:
            if isinstance(result, Exception):
                logger.error(f"Source query raised exception: {result}")
                continue
            if isinstance(result, list):
                merged.extend(result)

        return merged

    @staticmethod
    def _deduplicate(results: list[QueryResult]) -> list[QueryResult]:
        """基于内容哈希去重，保留分数最高的结果

        Args:
            results: 待去重的结果列表

        Returns:
            去重后的结果列表
        """
        seen: dict[str, QueryResult] = {}
        for r in results:
            content_key = _content_hash(r.content)
            if content_key not in seen or r.score > seen[content_key].score:
                seen[content_key] = r
        return list(seen.values())
