"""
MCP工具处理器（增强版）
使用依赖注入获取服务实例
支持后处理器、元数据过滤、批量查询和增量索引管理
"""
import json
from typing import Any, Dict, Optional, TYPE_CHECKING
from loguru import logger

from mcp.types import Tool, TextContent

if TYPE_CHECKING:
    from src.services import RAGService


class MCPToolsHandler:
    """
    MCP工具处理器（增强版）
    
    支持依赖注入，可以接受外部传入的 RAGService 实例。
    """
    
    def __init__(self, rag_service: Optional["RAGService"] = None):
        """
        初始化工具处理器
        
        Args:
            rag_service: RAGService 实例（可选）。
                        如果未提供，将从依赖注入容器获取。
        """
        if rag_service is None:
            from src.di import get_container
            container = get_container()
            rag_service = container.get("rag_service")
        
        self.service = rag_service
    
    @staticmethod
    def format_error(error: Exception) -> str:
        """格式化错误信息"""
        return f"错误: {type(error).__name__} - {str(error)}"
    
    def list_tools(self) -> list[Tool]:
        """列出所有可用的工具"""
        return [
            # ========== 现有工具 ==========
            Tool(
                name="build_index",
                description="构建知识库索引。从指定目录读取文档并构建向量索引和BM25索引。",
                inputSchema={
                    "type": "object",
                    "properties": {
                        "data_dir": {
                            "type": "string",
                            "description": "文档目录路径（绝对路径或相对路径）"
                        },
                        "force_rebuild": {
                            "type": "boolean",
                            "description": "是否强制重建索引（即使已存在），默认为false",
                            "default": False
                        }
                    },
                    "required": ["data_dir"]
                }
            ),
            Tool(
                name="query_knowledge",
                description="查询知识库并使用LLM生成答案。支持向量检索、BM25检索和混合检索三种模式。",
                inputSchema={
                    "type": "object",
                    "properties": {
                        "question": {
                            "type": "string",
                            "description": "用户问题"
                        },
                        "mode": {
                            "type": "string",
                            "description": "检索模式：vector（向量）、bm25（BM25）或hybrid（混合），默认为hybrid",
                            "enum": ["vector", "bm25", "hybrid"],
                            "default": "hybrid"
                        },
                        "top_k": {
                            "type": "integer",
                            "description": "检索的文档数量，默认为3",
                            "default": 3,
                            "minimum": 1,
                            "maximum": 20
                        }
                    },
                    "required": ["question"]
                }
            ),
            Tool(
                name="retrieve_documents",
                description="检索相关文档片段（增强版），支持相似度过滤、重排序和元数据过滤。不使用LLM生成答案。",
                inputSchema={
                    "type": "object",
                    "properties": {
                        "question": {
                            "type": "string",
                            "description": "检索内容或问题"
                        },
                        "mode": {
                            "type": "string",
                            "description": "检索模式：vector（向量）、bm25（BM25）或hybrid（混合），默认为hybrid",
                            "enum": ["vector", "bm25", "hybrid"],
                            "default": "hybrid"
                        },
                        "top_k": {
                            "type": "integer",
                            "description": "检索的文档数量，默认为5",
                            "default": 5,
                            "minimum": 1,
                            "maximum": 100
                        },
                        "similarity_cutoff": {
                            "type": "number",
                            "description": "相似度阈值（0-1），低于此值的结果将被过滤",
                            "minimum": 0.0,
                            "maximum": 1.0
                        },
                        "use_reranker": {
                            "type": "boolean",
                            "description": "是否使用重排序器提升结果质量"
                        },
                        "metadata_filters": {
                            "type": "object",
                            "description": "元数据过滤条件，格式: {\"key\": {\"operator\": \"EQ\", \"value\": \"xxx\"}}"
                        }
                    },
                    "required": ["question"]
                }
            ),
            Tool(
                name="get_index_status",
                description="获取索引状态和统计信息。返回向量索引和BM25索引的详细信息。",
                inputSchema={
                    "type": "object",
                    "properties": {
                        "verbose": {
                            "type": "boolean",
                            "description": "是否返回详细信息，默认为false",
                            "default": False
                        }
                    }
                }
            ),
            Tool(
                name="get_config",
                description="获取系统配置信息。返回所有配置参数，包括模型、API、检索参数等。",
                inputSchema={
                    "type": "object",
                    "properties": {}
                }
            ),
            Tool(
                name="delete_index",
                description="删除所有索引（包括向量索引和BM25索引）。警告：此操作不可逆！",
                inputSchema={
                    "type": "object",
                    "properties": {
                        "confirm": {
                            "type": "boolean",
                            "description": "确认删除操作（必须设置为true才能执行删除），默认为false",
                            "default": False
                        }
                    },
                    "required": []
                }
            ),
            
            # ========== 新增工具 ==========
            Tool(
                name="batch_retrieve",
                description="批量异步检索多个问题，高效处理多个查询请求。",
                inputSchema={
                    "type": "object",
                    "properties": {
                        "questions": {
                            "type": "array",
                            "items": {"type": "string"},
                            "description": "问题列表"
                        },
                        "mode": {
                            "type": "string",
                            "description": "检索模式",
                            "enum": ["vector", "bm25", "hybrid"],
                            "default": "hybrid"
                        },
                        "top_k": {
                            "type": "integer",
                            "description": "每个问题检索的文档数量",
                            "default": 5,
                            "minimum": 1,
                            "maximum": 20
                        }
                    },
                    "required": ["questions"]
                }
            ),
            Tool(
                name="insert_documents",
                description="插入新文档到索引（增量添加），无需重建整个索引。",
                inputSchema={
                    "type": "object",
                    "properties": {
                        "documents": {
                            "type": "array",
                            "items": {
                                "type": "object",
                                "properties": {
                                    "text": {"type": "string", "description": "文档内容"},
                                    "metadata": {"type": "object", "description": "文档元数据"}
                                },
                                "required": ["text"]
                            },
                            "description": "要插入的文档列表"
                        }
                    },
                    "required": ["documents"]
                }
            ),
            Tool(
                name="update_documents",
                description="更新索引中的已有文档。",
                inputSchema={
                    "type": "object",
                    "properties": {
                        "documents": {
                            "type": "array",
                            "items": {
                                "type": "object",
                                "properties": {
                                    "doc_id": {"type": "string", "description": "文档ID"},
                                    "text": {"type": "string", "description": "更新后的文档内容"},
                                    "metadata": {"type": "object", "description": "更新后的元数据"}
                                },
                                "required": ["doc_id", "text"]
                            },
                            "description": "要更新的文档列表"
                        }
                    },
                    "required": ["documents"]
                }
            ),
            Tool(
                name="delete_document",
                description="从索引中删除指定文档。",
                inputSchema={
                    "type": "object",
                    "properties": {
                        "doc_id": {
                            "type": "string",
                            "description": "要删除的文档ID"
                        }
                    },
                    "required": ["doc_id"]
                }
            ),
            Tool(
                name="list_documents",
                description="列出索引中的所有文档。",
                inputSchema={
                    "type": "object",
                    "properties": {
                        "limit": {
                            "type": "integer",
                            "description": "返回的文档数量限制",
                            "default": 100,
                            "minimum": 1,
                            "maximum": 1000
                        }
                    }
                }
            ),
            Tool(
                name="get_document",
                description="获取索引中指定文档的详细信息。",
                inputSchema={
                    "type": "object",
                    "properties": {
                        "doc_id": {
                            "type": "string",
                            "description": "文档ID"
                        }
                    },
                    "required": ["doc_id"]
                }
            )
        ]
    
    def call_tool(self, name: str, arguments: Dict[str, Any]) -> list[TextContent]:
        """处理工具调用"""
        try:
            logger.info(f"调用工具: {name}, 参数: {arguments}")
            
            # ========== 现有工具 ==========
            if name == "build_index":
                data_dir = arguments.get("data_dir")
                force_rebuild = arguments.get("force_rebuild", False)
                
                if not data_dir:
                    return [TextContent(
                        type="text",
                        text=json.dumps({
                            "success": False,
                            "error": "缺少必需参数: data_dir"
                        }, ensure_ascii=False, indent=2)
                    )]
                
                result = self.service.build_index(data_dir, force_rebuild)
                
                return [TextContent(
                    type="text",
                    text=json.dumps(result, ensure_ascii=False, indent=2)
                )]
            
            elif name == "query_knowledge":
                question = arguments.get("question")
                mode = arguments.get("mode")
                top_k = arguments.get("top_k")
                
                if not question:
                    return [TextContent(
                        type="text",
                        text=json.dumps({
                            "success": False,
                            "error": "缺少必需参数: question"
                        }, ensure_ascii=False, indent=2)
                    )]
                
                result = self.service.query(question=question, mode=mode, top_k=top_k)
                
                return [TextContent(
                    type="text",
                    text=json.dumps(result, ensure_ascii=False, indent=2)
                )]
            
            elif name == "retrieve_documents":
                question = arguments.get("question")
                mode = arguments.get("mode")
                top_k = arguments.get("top_k")
                similarity_cutoff = arguments.get("similarity_cutoff")
                use_reranker = arguments.get("use_reranker")
                metadata_filters = arguments.get("metadata_filters")
                
                if not question:
                    return [TextContent(
                        type="text",
                        text=json.dumps({
                            "success": False,
                            "error": "缺少必需参数: question"
                        }, ensure_ascii=False, indent=2)
                    )]
                
                result = self.service.retrieve(
                    question=question, 
                    mode=mode, 
                    top_k=top_k,
                    similarity_cutoff=similarity_cutoff,
                    use_reranker=use_reranker,
                    metadata_filters=metadata_filters
                )
                
                return [TextContent(
                    type="text",
                    text=json.dumps(result, ensure_ascii=False, indent=2)
                )]
            
            elif name == "get_index_status":
                verbose = arguments.get("verbose", False)
                
                result = self.service.get_status(verbose=verbose)
                
                return [TextContent(
                    type="text",
                    text=json.dumps(result, ensure_ascii=False, indent=2)
                )]
            
            elif name == "get_config":
                result = self.service.get_config()
                
                return [TextContent(
                    type="text",
                    text=json.dumps(result, ensure_ascii=False, indent=2)
                )]
            
            elif name == "delete_index":
                confirm = arguments.get("confirm", False)
                
                result = self.service.delete_index(confirm=confirm)
                
                return [TextContent(
                    type="text",
                    text=json.dumps(result, ensure_ascii=False, indent=2)
                )]
            
            # ========== 新增工具 ==========
            elif name == "batch_retrieve":
                questions = arguments.get("questions")
                mode = arguments.get("mode")
                top_k = arguments.get("top_k")
                
                if not questions or not isinstance(questions, list):
                    return [TextContent(
                        type="text",
                        text=json.dumps({
                            "success": False,
                            "error": "缺少必需参数: questions (必须是数组)"
                        }, ensure_ascii=False, indent=2)
                    )]
                
                result = self.service.batch_retrieve(
                    questions=questions,
                    mode=mode,
                    top_k=top_k
                )
                
                return [TextContent(
                    type="text",
                    text=json.dumps(result, ensure_ascii=False, indent=2)
                )]
            
            elif name == "insert_documents":
                documents = arguments.get("documents")
                
                if not documents or not isinstance(documents, list):
                    return [TextContent(
                        type="text",
                        text=json.dumps({
                            "success": False,
                            "error": "缺少必需参数: documents (必须是数组)"
                        }, ensure_ascii=False, indent=2)
                    )]
                
                result = self.service.insert_documents(documents)
                
                return [TextContent(
                    type="text",
                    text=json.dumps(result, ensure_ascii=False, indent=2)
                )]
            
            elif name == "update_documents":
                documents = arguments.get("documents")
                
                if not documents or not isinstance(documents, list):
                    return [TextContent(
                        type="text",
                        text=json.dumps({
                            "success": False,
                            "error": "缺少必需参数: documents (必须是数组)"
                        }, ensure_ascii=False, indent=2)
                    )]
                
                result = self.service.update_documents(documents)
                
                return [TextContent(
                    type="text",
                    text=json.dumps(result, ensure_ascii=False, indent=2)
                )]
            
            elif name == "delete_document":
                doc_id = arguments.get("doc_id")
                
                if not doc_id:
                    return [TextContent(
                        type="text",
                        text=json.dumps({
                            "success": False,
                            "error": "缺少必需参数: doc_id"
                        }, ensure_ascii=False, indent=2)
                    )]
                
                result = self.service.delete_document(doc_id)
                
                return [TextContent(
                    type="text",
                    text=json.dumps(result, ensure_ascii=False, indent=2)
                )]
            
            elif name == "list_documents":
                limit = arguments.get("limit", 100)
                
                result = self.service.list_documents(limit=limit)
                
                return [TextContent(
                    type="text",
                    text=json.dumps(result, ensure_ascii=False, indent=2)
                )]
            
            elif name == "get_document":
                doc_id = arguments.get("doc_id")
                
                if not doc_id:
                    return [TextContent(
                        type="text",
                        text=json.dumps({
                            "success": False,
                            "error": "缺少必需参数: doc_id"
                        }, ensure_ascii=False, indent=2)
                    )]
                
                result = self.service.get_document(doc_id)
                
                return [TextContent(
                    type="text",
                    text=json.dumps(result, ensure_ascii=False, indent=2)
                )]
            
            else:
                return [TextContent(
                    type="text",
                    text=json.dumps({
                        "success": False,
                        "error": f"未知工具: {name}"
                    }, ensure_ascii=False, indent=2)
                )]
        
        except Exception as e:
            logger.error(f"工具调用失败: {e}")
            import traceback
            logger.error(traceback.format_exc())
            
            return [TextContent(
                type="text",
                text=json.dumps({
                    "success": False,
                    "error": self.format_error(e),
                    "traceback": traceback.format_exc()
                }, ensure_ascii=False, indent=2)
            )]
