"""知识检索工具

KnowledgeRetrievalTool - 从知识库检索解构所需的知识
"""

from typing import Dict, Any, List, Optional
from pydantic import BaseModel, Field
from langchain_core.tools import BaseTool
import requests
import urllib.parse

from src.components.tools.base import SimpleTool, AsyncTool
from src.utils.logger import get_logger

logger = get_logger(__name__)


class KnowledgeRetrievalInput(BaseModel):
    """知识检索工具输入Schema"""
    query: str = Field(..., description="知识查询文本")
    knowledge_type: str = Field(
        default="general",
        description="知识类型: category(品类知识), element(元素知识), split(拆分知识), general(通用)"
    )
    top_k: int = Field(default=3, description="返回Top K结果", ge=1, le=10)


class KnowledgeRetrievalTool:
    """知识检索工具

    根据Query从知识库检索解构所需的知识

    功能:
    - 接收Query并调用知识库API
    - 支持多种知识库类型（向量检索、关系型等）
    - 返回结构化知识

    实现方式: Tool组件（确定性API调用 + 需要被Agent调用）
    依据: PRD 3.3节"外部知识调用"
    """

    def __init__(
        self,
        api_url: str = "http://61.48.133.26:8001/api/query",
        dataset_ids: str = "11",
        timeout: int = 10
    ):
        """初始化知识检索工具

        Args:
            api_url: 知识库API地址
            dataset_ids: 数据集ID，逗号分隔（11=内容知识库V2, 12=工具使用库V2）
            timeout: API请求超时时间（秒）
        """
        self.api_url = api_url
        self.dataset_ids = dataset_ids
        self.timeout = timeout
        self._tool: Optional[SimpleTool] = None

    def _retrieve_knowledge_api(
        self,
        query: str,
        knowledge_type: str = "general",
        top_k: int = 3
    ) -> Dict[str, Any]:
        """通过API检索知识

        Args:
            query: 查询文本
            knowledge_type: 知识类型（仅用于标记，API统一检索）
            top_k: 返回Top K结果

        Returns:
            检索结果字典
        """
        logger.info(f"API知识检索 - Query: {query}, Type: {knowledge_type}, Top_k: {top_k}")

        try:
            # 构建请求URL
            params = {
                "query": query,
                "datasetIds": self.dataset_ids
            }
            url = f"{self.api_url}?{urllib.parse.urlencode(params)}"

            # 发送请求
            logger.debug(f"请求URL: {url}")
            response = requests.get(url, timeout=self.timeout)
            response.raise_for_status()

            # 解析响应
            data = response.json()

            if data.get("status_code") != 200:
                logger.warning(f"API返回错误: {data.get('detail', 'Unknown error')}")
                return self._create_empty_result(query, knowledge_type, "API返回错误")

            # 提取检索结果
            results = data.get("data", {}).get("results", [])

            if not results:
                logger.info("API未返回任何结果，使用默认知识")

            # 限制返回数量
            results = results[:top_k]

            # 整理知识内容
            knowledge_items = []
            for item in results:
                content = item.get("content", item.get("content", ""))
                if content:
                    knowledge_items.append({
                        "content": content,
                        "source": item.get("datasetName", "Unknown"),
                        "score": item.get("score", 0.0),
                        "doc_id": item.get("docId", "")
                    })

            # 合并知识内容
            combined_knowledge = "\n\n".join([
                f"[来源: {item['source']}]\n{item['content']}"
                for item in knowledge_items
            ])

            # 计算平均置信度
            avg_confidence = sum([item["score"] for item in knowledge_items]) / len(knowledge_items) if knowledge_items else 0.0

            result = {
                "knowledge": combined_knowledge,
                "sources": [item["source"] for item in knowledge_items],
                "confidence": avg_confidence,
                "query": query,
                "knowledge_type": knowledge_type,
                "results_count": len(knowledge_items),
                "api_response": True
            }

            logger.info(f"API检索成功，返回 {len(knowledge_items)} 条知识，平均置信度: {avg_confidence:.4f}")
            logger.info(f"API检索成功，返回 - Query: {query}, 知识: {combined_knowledge}")

            return result

        except requests.exceptions.Timeout:
            logger.error(f"API请求超时（{self.timeout}秒）")
            return self._create_empty_result(query, knowledge_type, "API请求超时，使用备用知识")

        except requests.exceptions.RequestException as e:
            logger.error(f"API请求失败: {e}")

        except Exception as e:
            logger.error(f"知识检索异常: {e}", exc_info=True)
            return self._create_empty_result(query, knowledge_type, f"检索异常: {str(e)}")

    def _create_empty_result(self, query: str, knowledge_type: str, reason: str) -> Dict[str, Any]:
        """创建空结果

        Args:
            query: 查询文本
            knowledge_type: 知识类型
            reason: 原因

        Returns:
            空结果字典
        """
        return {
            "knowledge": f"暂无相关知识。原因: {reason}",
            "sources": [],
            "confidence": 0.0,
            "query": query,
            "knowledge_type": knowledge_type,
            "results_count": 0,
            "error": reason
        }

    def _create_tool_function(self):
        """创建工具函数"""
        def retrieve_knowledge(
            query: str,
            knowledge_type: str = "general",
            top_k: int = 3
        ) -> Dict[str, Any]:
            """从知识库检索解构所需的知识

            Args:
                query: 知识查询文本
                knowledge_type: 知识类型 (category/element/split/general)
                top_k: 返回Top K结果

            Returns:
                包含知识内容、来源和置信度的字典
            """
            return self._retrieve_knowledge_api(query, knowledge_type, top_k)

        return retrieve_knowledge

    def get_tool(self) -> BaseTool:
        """获取工具实例

        Returns:
            LangChain BaseTool实例
        """
        if self._tool is None:
            func = self._create_tool_function()
            self._tool = SimpleTool(
                name="knowledge_retrieval",
                description="从知识库检索What解构所需的知识。支持品类知识、元素知识、拆分知识等多种类型的查询。",
                func=func,
                schema=KnowledgeRetrievalInput
            )
            self._tool.initialize()

        return self._tool.get_tool()

    def retrieve(
        self,
        query: str,
        knowledge_type: str = "general",
        top_k: int = 3
    ) -> Dict[str, Any]:
        """直接调用检索方法（非Tool调用）

        Args:
            query: 查询文本
            knowledge_type: 知识类型
            top_k: 返回Top K结果

        Returns:
            检索结果
        """
        func = self._create_tool_function()
        return func(query=query, knowledge_type=knowledge_type, top_k=top_k)


# 创建默认实例（使用真实API）
default_knowledge_tool = KnowledgeRetrievalTool()