class RoutingRAG:
    """路由式 RAG 架构"""
    
    def __init__(self):
        self.query_classifier = QueryClassifier()
        self.retrieval_strategies = {
            'factual': FactualRetriever(),      # 事实查询
            'analytical': AnalyticalRetriever(), # 分析查询
            'comparative': ComparativeRetriever(), # 对比查询
            'procedural': ProceduralRetriever()  # 流程查询
        }
    
    async def process(self, query: str) -> dict:
        """处理查询"""
        # 1. 分类查询类型
        query_type = await self.query_classifier.classify(query)
        
        # 2. 路由到对应策略
        retriever = self.retrieval_strategies[query_type]
        
        # 3. 执行检索
        results = await retriever.retrieve(query)
        
        # 4. 生成答案
        answer = await self._generate_answer(query, results, query_type)
        
        return {
            'answer': answer,
            'query_type': query_type,
            'sources': results
        }

class QueryClassifier:
    """查询分类器"""
    
    async def classify(self, query: str) -> str:
        """
        分类规则：
        - 事实查询: "什么是...", "谁是...", "何时..."
        - 分析查询: "为什么...", "如何解释...", "分析..."
        - 对比查询: "A和B的区别", "比较...", "哪个更好"
        - 流程查询: "如何做...", "步骤...", "流程..."
        """
        prompt = f"""
        分类以下查询的类型：
        查询: {query}
        
        类型选项:
        1. factual - 事实性查询
        2. analytical - 分析性查询
        3. comparative - 对比性查询
        4. procedural - 流程性查询
        
        只返回类型名称。
        """
        
        response = await self.llm.acomplete(prompt)
        return response.text.strip().lower()

class QueryRewritingRAG:
    """查询改写 RAG"""
    
    async def process(self, query: str) -> dict:
        """处理流程"""
        
        # 1. 生成多个查询变体
        rewritten_queries = await self._rewrite_query(query)
        
        # 2. 对每个变体进行检索
        all_results = []
        for q in [query] + rewritten_queries:
            results = await self.retriever.retrieve(q)
            all_results.extend(results)
        
        # 3. 去重和重排序
        unique_results = self._deduplicate(all_results)
        reranked = await self.reranker.rerank(query, unique_results)
        
        # 4. 生成答案
        answer = await self._generate(query, reranked[:5])
        
        return {
            'answer': answer,
            'rewritten_queries': rewritten_queries,
            'sources': reranked[:5]
        }
    
    async def _rewrite_query(self, query: str) -> List[str]:
        """查询改写策略"""
        strategies = [
            self._expand_with_synonyms,      # 同义词扩展
            self._simplify_query,            # 简化复杂查询
            self._add_context,               # 添加上下文
            self._decompose_complex_query    # 分解复杂查询
        ]
        
        rewritten = []
        for strategy in strategies:
            variant = await strategy(query)
            if variant and variant != query:
                rewritten.append(variant)
        
        return rewritten[:3]  # 最多3个变体
    
    async def _expand_with_synonyms(self, query: str) -> str:
        """使用同义词扩展"""
        prompt = f"""
        扩展以下查询，使用同义词替换关键词：
        原查询: {query}
        
        返回一个语义相同但用词不同的查询。
        """
        response = await self.llm.acomplete(prompt)
        return response.text.strip()
    
    async def _decompose_complex_query(self, query: str) -> str:
        """分解复杂查询"""
        prompt = f"""
        将以下复杂查询分解为一个更简单、更聚焦的子查询：
        查询: {query}
        
        只返回最核心的子查询。
        """
        response = await self.llm.acomplete(prompt)
        return response.text.strip()

class AdaptiveRAG:
    """自适应 RAG - 根据复杂度选择策略"""
    
    async def process(self, query: str) -> dict:
        """自适应处理"""
        
        # 1. 评估查询复杂度
        complexity = await self._assess_complexity(query)
        
        # 2. 根据复杂度选择策略
        if complexity == 'simple':
            # 简单查询：单次检索
            return await self._simple_retrieval(query)
        
        elif complexity == 'medium':
            # 中等复杂度：查询改写 + 重排序
            return await self._enhanced_retrieval(query)
        
        else:
            # 高复杂度：多跳推理 + Agent协作
            return await self._complex_reasoning(query)
    
    async def _assess_complexity(self, query: str) -> str:
        """评估查询复杂度"""
        factors = {
            'token_count': len(self.tokenizer.encode(query)),
            'question_count': query.count('?'),
            'has_comparison': any(kw in query.lower() for kw in ['比较', '区别', '对比', 'vs']),
            'has_reasoning': any(kw in query.lower() for kw in ['为什么', '如何', '解释', '分析']),
            'has_multi_steps': any(kw in query.lower() for kw in ['首先', '然后', '步骤', '流程'])
        }
        
        score = 0
        if factors['token_count'] > 50: score += 2
        if factors['question_count'] > 1: score += 2
        if factors['has_comparison']: score += 1
        if factors['has_reasoning']: score += 1
        if factors['has_multi_steps']: score += 2
        
        if score <= 2:
            return 'simple'
        elif score <= 5:
            return 'medium'
        else:
            return 'complex'
    
    async def _complex_reasoning(self, query: str) -> dict:
        """复杂推理流程"""
        # 使用 ReAct 模式
        state = {
            'query': query,
            'observations': [],
            'reasoning_steps': []
        }
        
        max_iterations = 5
        for i in range(max_iterations):
            # Reasoning: 分析当前情况
            thought = await self._reason(state)
            state['reasoning_steps'].append(thought)
            
            # Action: 执行检索或工具调用
            action_result = await self._act(thought)
            
            # Observation: 观察结果
            state['observations'].append(action_result)
            
            # Reflection: 判断是否需要继续
            is_complete = await self._reflect(state)
            if is_complete:
                break
        
        # 最终答案生成
        answer = await self._synthesize(state)
        
        return {
            'answer': answer,
            'reasoning_steps': state['reasoning_steps'],
            'observations': state['observations']
        }

class CorrectiveRAG:
    """纠错式 RAG"""
    
    async def process(self, query: str) -> dict:
        """处理流程"""
        
        # 1. 初始检索
        initial_results = await self.retriever.retrieve(query, top_k=10)
        
        # 2. 评估检索质量
        quality_scores = await self._evaluate_quality(query, initial_results)
        
        # 3. 决策：使用、改进或放弃
        decision = self._make_decision(quality_scores)
        
        if decision == 'use':
            # 检索结果质量好，直接使用
            final_results = initial_results[:5]
        
        elif decision == 'refine':
            # 检索结果部分相关，需要改进
            final_results = await self._refine_retrieval(
                query, 
                initial_results,
                quality_scores
            )
        
        else:  # decision == 'search_web'
            # 检索结果不相关，使用网络搜索
            final_results = await self._web_search(query)
        
        # 4. 生成答案
        answer = await self._generate(query, final_results)
        
        return {
            'answer': answer,
            'decision': decision,
            'quality_scores': quality_scores,
            'sources': final_results
        }
    
    async def _evaluate_quality(
        self, 
        query: str, 
        results: List
    ) -> List[float]:
        """评估检索质量"""
        scores = []
        
        for result in results:
            # 使用 LLM 评估相关性
            prompt = f"""
            评估以下文档与查询的相关性（0-10分）：
            
            查询: {query}
            
            文档: {result.get_content()[:500]}
            
            只返回分数（整数）。
            """
            
            response = await self.llm.acomplete(prompt)
            score = int(response.text.strip())
            scores.append(score / 10.0)  # 归一化到 0-1
        
        return scores
    
    def _make_decision(self, scores: List[float]) -> str:
        """决策逻辑"""
        avg_score = sum(scores) / len(scores) if scores else 0
        max_score = max(scores) if scores else 0
        
        if avg_score >= 0.7:
            return 'use'  # 平均分高，直接使用
        elif max_score >= 0.6:
            return 'refine'  # 有部分相关文档，改进检索
        else:
            return 'search_web'  # 都不相关，使用网络搜索
    
    async def _refine_retrieval(
        self,
        query: str,
        initial_results: List,
        scores: List[float]
    ) -> List:
        """改进检索"""
        # 1. 从高分文档中提取关键概念
        good_docs = [
            doc for doc, score in zip(initial_results, scores)
            if score >= 0.5
        ]
        
        key_concepts = await self._extract_concepts(good_docs)
        
        # 2. 使用关键概念扩展查询
        expanded_query = f"{query} {' '.join(key_concepts)}"
        
        # 3. 重新检索
        refined_results = await self.retriever.retrieve(
            expanded_query, 
            top_k=5
        )
        
        return refined_results

class SelfRAG:
    """自省式 RAG"""
    
    async def process(self, query: str) -> dict:
        """处理流程"""
        
        response_parts = []
        reflections = []
        
        # 1. 初始判断：是否需要检索？
        need_retrieval = await self._should_retrieve(query)
        
        if not need_retrieval:
            # 直接使用参数化知识回答
            answer = await self.llm.acomplete(query)
            return {
                'answer': answer.text,
                'retrieval_used': False
            }
        
        # 2. 执行检索
        retrieved_docs = await self.retriever.retrieve(query)
        
        # 3. 逐段生成，并自我评估
        context = "\n\n".join([doc.get_content() for doc in retrieved_docs])
        
        # 生成第一段
        segment = await self._generate_segment(query, context)
        response_parts.append(segment)
        
        # 4. 自我反思：是否需要继续？
        for i in range(3):  # 最多3轮
            reflection = await self._reflect(query, response_parts, context)
            reflections.append(reflection)
            
            if reflection['is_complete']:
                break
            
            if reflection['need_more_info']:
                # 需要更多信息，再次检索
                new_query = reflection['refined_query']
                new_docs = await self.retriever.retrieve(new_query)
                context += "\n\n" + "\n\n".join([
                    doc.get_content() for doc in new_docs
                ])
            
            # 生成下一段
            next_segment = await self._generate_segment(
                query, 
                context,
                previous=response_parts
            )
            response_parts.append(next_segment)
        
        # 5. 最终验证
        final_answer = " ".join(response_parts)
        is_supported = await self._verify_support(final_answer, retrieved_docs)
        
        return {
            'answer': final_answer,
            'retrieval_used': True,
            'is_supported': is_supported,
            'reflections': reflections
        }
    
    async def _should_retrieve(self, query: str) -> bool:
        """判断是否需要检索"""
        prompt = f"""
        判断以下问题是否需要外部知识：
        
        问题: {query}
        
        如果问题关于事实、数据、具体信息，返回 YES
        如果问题关于常识、推理、创意，返回 NO
        
        只返回 YES 或 NO。
        """
        
        response = await self.llm.acomplete(prompt)
        return 'YES' in response.text.upper()
    
    async def _reflect(
        self,
        query: str,
        generated: List[str],
        context: str
    ) -> dict:
        """自我反思"""
        current_answer = " ".join(generated)
        
        prompt = f"""
        评估当前答案的完整性：
        
        问题: {query}
        当前答案: {current_answer}
        可用上下文: {context[:1000]}...
        
        回答以下问题：
        1. 答案是否完整？(YES/NO)
        2. 是否需要更多信息？(YES/NO)
        3. 如果需要，应该搜索什么？
        
        返回JSON格式。
        """
        
        response = await self.llm.acomplete(prompt)
        
        # 解析响应
        reflection = json.loads(response.text)
        
        return {
            'is_complete': reflection.get('is_complete', False),
            'need_more_info': reflection.get('need_more_info', False),
            'refined_query': reflection.get('refined_query', '')
        }

class GraphRAG:
    """图谱增强 RAG"""
    
    def __init__(
        self,
        vector_store,
        knowledge_graph,  # Neo4j 或其他图数据库
    ):
        self.vector_store = vector_store
        self.kg = knowledge_graph
    
    async def process(self, query: str) -> dict:
        """处理流程"""
        
        # 1. 实体识别
        entities = await self._extract_entities(query)
        
        # 2. 并行执行向量检索和图谱检索
        vector_task = self._vector_retrieve(query)
        graph_task = self._graph_retrieve(entities)
        
        vector_results, graph_results = await asyncio.gather(
            vector_task,
            graph_task
        )
        
        # 3. 融合结果
        fused_context = self._fuse_contexts(
            vector_results,
            graph_results
        )
        
        # 4. 生成答案
        answer = await self._generate(query, fused_context)
        
        return {
            'answer': answer,
            'entities': entities,
            'vector_sources': vector_results,
            'graph_relations': graph_results
        }
    
    async def _graph_retrieve(self, entities: List[str]) -> dict:
        """从知识图谱检索"""
        graph_context = {}
        
        for entity in entities:
            # Cypher 查询：获取实体的关系
            cypher = f"""
            MATCH (e:Entity {{name: $entity}})-[r]->(related)
            RETURN e, type(r) as relation, related
            LIMIT 10
            """
            
            results = await self.kg.execute(cypher, entity=entity)
            
            graph_context[entity] = {
                'direct_relations': results,
                'properties': await self._get_entity_properties(entity)
            }
        
        return graph_context
    
    def _fuse_contexts(
        self,
        vector_results: List,
        graph_results: dict
    ) -> str:
        """融合向量检索和图谱结果"""
        
        # 向量检索的文本内容
        text_context = "\n\n".join([
            f"文档 {i+1}:\n{doc.get_content()}"
            for i, doc in enumerate(vector_results)
        ])
        
        # 图谱的结构化信息
        graph_context = "相关实体和关系:\n"
        for entity, info in graph_results.items():
            graph_context += f"\n{entity}:\n"
            for rel in info['direct_relations']:
                graph_context += f"  - {rel['relation']}: {rel['related']['name']}\n"
        
        return f"{text_context}\n\n{graph_context}"

class MultimodalRAG:
    """多模态 RAG"""
    
    def __init__(self):
        self.text_retriever = VectorRetriever()
        self.image_retriever = CLIPRetriever()
        self.table_retriever = TableRetriever()
    
    async def process(self, query: str, query_image=None) -> dict:
        """处理多模态查询"""
        
        # 1. 识别查询涉及的模态
        modalities = await self._detect_modalities(query)
        
        # 2. 多模态检索
        results = {}
        
        if 'text' in modalities:
            results['text'] = await self.text_retriever.retrieve(query)
        
        if 'image' in modalities or query_image:
            results['images'] = await self.image_retriever.retrieve(
                query, 
                query_image
            )
        
        if 'table' in modalities:
            results['tables'] = await self.table_retriever.retrieve(query)
        
        # 3. 多模态融合
        fused_context = await self._fuse_multimodal(results)
        
        # 4. 使用多模态模型生成答案
        answer = await self._multimodal_generate(
            query,
            fused_context,
            query_image
        )
        
        return {
            'answer': answer,
            'modalities_used': list(results.keys()),
            'sources': results
        }
    
    async def _detect_modalities(self, query: str) -> List[str]:
        """检测查询涉及的模态"""
        modalities = ['text']  # 默认包含文本
        
        # 关键词检测
        if any(kw in query.lower() for kw in ['图片', '图像', '照片', 'image', 'picture']):
            modalities.append('image')
        
        if any(kw in query.lower() for kw in ['表格', '数据', 'table', 'chart']):
            modalities.append('table')
        
        return modalities
    
    async def _fuse_multimodal(self, results: dict) -> dict:
        """融合多模态结果"""
        fused = {}
        
        if 'text' in results:
            fused['text_context'] = "\n\n".join([
                doc.get_content() for doc in results['text']
            ])
        
        if 'images' in results:
            # 使用 CLIP 提取图像特征和描述
            fused['image_descriptions'] = [
                await self._describe_image(img) 
                for img in results['images']
            ]
        
        if 'tables' in results:
            # 将表格转换为结构化文本
            fused['table_data'] = [
                self._table_to_text(table)
                for table in results['tables']
            ]
        
        return fused

from mcp.server import Server
from mcp.server.stdio import stdio_server
from mcp.types import Tool, TextContent
from typing import Any
import asyncio

# ============= MCP Server 基类 =============
class BaseMCPServer:
    """MCP Server 基类"""
    
    def __init__(self, name: str):
        self.name = name
        self.server = Server(name)
        self._setup_handlers()
    
    def _setup_handlers(self):
        """设置处理器"""
        
        @self.server.list_tools()
        async def list_tools() -> list[Tool]:
            """列出可用工具"""
            return await self.get_tools()
        
        @self.server.call_tool()
        async def call_tool(name: str, arguments: dict) -> list[TextContent]:
            """调用工具"""
            result = await self.execute_tool(name, arguments)
            return [TextContent(
                type="text",
                text=str(result)
            )]
    
    async def get_tools(self) -> list[Tool]:
        """子类实现：返回工具列表"""
        raise NotImplementedError
    
    async def execute_tool(self, name: str, arguments: dict) -> Any:
        """子类实现：执行工具"""
        raise NotImplementedError
    
    async def run(self):
        """运行 MCP Server"""
        async with stdio_server() as (read_stream, write_stream):
            await self.server.run(
                read_stream,
                write_stream,
                self.server.create_initialization_options()
            )

# ============= S3 数据源 MCP Server =============
class S3MCPServer(BaseMCPServer):
    """S3 数据源 MCP Server"""
    
    def __init__(self, aws_access_key: str, aws_secret_key: str):
        super().__init__("s3-datasource")
        self.s3_client = boto3.client(
            's3',
            aws_access_key_id=aws_access_key,
            aws_secret_access_key=aws_secret_key
        )
    
    async def get_tools(self) -> list[Tool]:
        """返回 S3 相关工具"""
        return [
            Tool(
                name="s3_list_objects",
                description="List objects in S3 bucket with optional prefix filter",
                inputSchema={
                    "type": "object",
                    "properties": {
                        "bucket": {
                            "type": "string",
                            "description": "S3 bucket name"
                        },
                        "prefix": {
                            "type": "string",
                            "description": "Optional prefix to filter objects"
                        },
                        "max_keys": {
                            "type": "integer",
                            "description": "Maximum number of keys to return",
                            "default": 1000
                        }
                    },
                    "required": ["bucket"]
                }
            ),
            Tool(
                name="s3_read_object",
                description="Read content from S3 object",
                inputSchema={
                    "type": "object",
                    "properties": {
                        "bucket": {"type": "string"},
                        "key": {"type": "string"},
                        "encoding": {
                            "type": "string",
                            "enum": ["utf-8", "binary"],
                            "default": "utf-8"
                        }
                    },
                    "required": ["bucket", "key"]
                }
            ),
            Tool(
                name="s3_sync_to_vectordb",
                description="Sync S3 objects to vector database",
                inputSchema={
                    "type": "object",
                    "properties": {
                        "bucket": {"type": "string"},
                        "prefix": {"type": "string"},
                        "mode": {
                            "type": "string",
                            "enum": ["full", "incremental"],
                            "default": "incremental"
                        }
                    },
                    "required": ["bucket"]
                }
            )
        ]
    
    async def execute_tool(self, name: str, arguments: dict) -> Any:
        """执行 S3 工具"""
        if name == "s3_list_objects":
            return await self._list_objects(**arguments)
        elif name == "s3_read_object":
            return await self._read_object(**arguments)
        elif name == "s3_sync_to_vectordb":
            return await self._sync_to_vectordb(**arguments)
        else:
            raise ValueError(f"Unknown tool: {name}")
    
    async def _list_objects(
        self, 
        bucket: str, 
        prefix: str = "",
        max_keys: int = 1000
    ) -> dict:
        """列出 S3 对象"""
        response = await asyncio.to_thread(
            self.s3_client.list_objects_v2,
            Bucket=bucket,
            Prefix=prefix,
            MaxKeys=max_keys
        )
        
        objects = []
        for obj in response.get('Contents', []):
            objects.append({
                'key': obj['Key'],
                'size': obj['Size'],
                'last_modified': obj['LastModified'].isoformat()
            })
        
        return {
            'objects': objects,
            'count': len(objects),
            'is_truncated': response.get('IsTruncated', False)
        }
    
    async def _read_object(
        self, 
        bucket: str, 
        key: str,
        encoding: str = "utf-8"
    ) -> str:
        """读取 S3 对象内容"""
        response = await asyncio.to_thread(
            self.s3_client.get_object,
            Bucket=bucket,
            Key=key
        )
        
        content = response['Body'].read()
        
        if encoding == "utf-8":
            return content.decode('utf-8')
        return content.hex()  # 返回十六进制字符串
    
    async def _sync_to_vectordb(
        self,
        bucket: str,
        prefix: str = "",
        mode: str = "incremental"
    ) -> dict:
        """同步到向量数据库"""
        # 这里会调用数据摄取层
        from data_ingestion import DataIngestionPipeline
        
        pipeline = DataIngestionPipeline()
        result = await pipeline.sync_s3_source(
            bucket=bucket,
            prefix=prefix,
            mode=mode
        )
        
        return result

# ============= Google Drive MCP Server =============
class GoogleDriveMCPServer(BaseMCPServer):
    """Google Drive MCP Server"""
    
    def __init__(self, credentials_path: str):
        super().__init__("google-drive-datasource")
        self.drive_service = self._init_drive_service(credentials_path)
    
    def _init_drive_service(self, credentials_path: str):
        """初始化 Google Drive 服务"""
        from google.oauth2.credentials import Credentials
        from googleapiclient.discovery import build
        
        creds = Credentials.from_authorized_user_file(
            credentials_path,
            scopes=['https://www.googleapis.com/auth/drive.readonly']
        )
        
        return build('drive', 'v3', credentials=creds)
    
    async def get_tools(self) -> list[Tool]:
        """返回 Google Drive 工具"""
        return [
            Tool(
                name="gdrive_list_files",
                description="List files in Google Drive",
                inputSchema={
                    "type": "object",
                    "properties": {
                        "folder_id": {
                            "type": "string",
                            "description": "Folder ID to list files from"
                        },
                        "query": {
                            "type": "string",
                            "description": "Search query"
                        },
                        "page_size": {
                            "type": "integer",
                            "default": 100
                        }
                    }
                }
            ),
            Tool(
                name="gdrive_read_file",
                description="Read file content from Google Drive",
                inputSchema={
                    "type": "object",
                    "properties": {
                        "file_id": {"type": "string"}
                    },
                    "required": ["file_id"]
                }
            ),
            Tool(
                name="gdrive_sync_to_vectordb",
                description="Sync Google Drive files to vector database",
                inputSchema={
                    "type": "object",
                    "properties": {
                        "folder_id": {"type": "string"},
                        "mode": {
                            "type": "string",
                            "enum": ["full", "incremental"]
                        }
                    }
                }
            )
        ]
    
    async def execute_tool(self, name: str, arguments: dict) -> Any:
        """执行工具"""
        if name == "gdrive_list_files":
            return await self._list_files(**arguments)
        elif name == "gdrive_read_file":
            return await self._read_file(**arguments)
        elif name == "gdrive_sync_to_vectordb":
            return await self._sync_to_vectordb(**arguments)
        else:
            raise ValueError(f"Unknown tool: {name}")
    
    async def _list_files(
        self,
        folder_id: str = None,
        query: str = None,
        page_size: int = 100
    ) -> dict:
        """列出文件"""
        q_parts = []
        if folder_id:
            q_parts.append(f"'{folder_id}' in parents")
        if query:
            q_parts.append(query)
        
        q_string = " and ".join(q_parts) if q_parts else None
        
        results = await asyncio.to_thread(
            self.drive_service.files().list,
            q=q_string,
            pageSize=page_size,
            fields="files(id, name, mimeType, modifiedTime, size)"
        ).execute()
        
        return {
            'files': results.get('files', []),
            'count': len(results.get('files', []))
        }
    
    async def _read_file(self, file_id: str) -> str:
        """读取文件内容"""
        # 获取文件元数据
        file_metadata = await asyncio.to_thread(
            self.drive_service.files().get,
            fileId=file_id,
            fields='mimeType'
        ).execute()
        
        mime_type = file_metadata['mimeType']
        
        # 根据 MIME 类型选择导出格式
        if 'google-apps' in mime_type:
            # Google Docs/Sheets/Slides 需要导出
            export_mime = 'text/plain'
            content = await asyncio.to_thread(
                self.drive_service.files().export_media,
                fileId=file_id,
                mimeType=export_mime
            ).execute()
        else:
            # 普通文件直接下载
            content = await asyncio.to_thread(
                self.drive_service.files().get_media,
                fileId=file_id
            ).execute()
        
        return content.decode('utf-8')

# ============= MCP 服务注册中心 =============
class MCPRegistry:
    """MCP Server 注册中心"""
    
    def __init__(self):
        self.servers: Dict[str, BaseMCPServer] = {}
        self.tools_cache: Dict[str, list[Tool]] = {}
    
    def register_server(self, server: BaseMCPServer):
        """注册 MCP Server"""
        self.servers[server.name] = server
        logger.info(f"Registered MCP server: {server.name}")
    
    async def discover_tools(self) -> Dict[str, list[Tool]]:
        """发现所有工具"""
        all_tools = {}
        
        for name, server in self.servers.items():
            tools = await server.get_tools()
            all_tools[name] = tools
            self.tools_cache[name] = tools
        
        return all_tools
    
    async def execute_tool(
        self,
        server_name: str,
        tool_name: str,
        arguments: dict
    ) -> Any:
        """执行工具"""
        if server_name not in self.servers:
            raise ValueError(f"Server {server_name} not found")
        
        server = self.servers[server_name]
        return await server.execute_tool(tool_name, arguments)
    
    def get_tool_by_name(self, tool_name: str) -> tuple[str, Tool]:
        """根据工具名称查找工具"""
        for server_name, tools in self.tools_cache.items():
            for tool in tools:
                if tool.name == tool_name:
                    return server_name, tool
        
        raise ValueError(f"Tool {tool_name} not found")

# ============= 使用示例 =============
async def setup_mcp_servers():
    """设置 MCP Servers"""
    registry = MCPRegistry()
    
    # 注册 S3 Server
    s3_server = S3MCPServer(
        aws_access_key=os.getenv('AWS_ACCESS_KEY'),
        aws_secret_key=os.getenv('AWS_SECRET_KEY')
    )
    registry.register_server(s3_server)
    
    # 注册 Google Drive Server
    gdrive_server = GoogleDriveMCPServer(
        credentials_path='./credentials.json'
    )
    registry.register_server(gdrive_server)
    
    # 发现所有工具
    tools = await registry.discover_tools()
    logger.info(f"Discovered {sum(len(t) for t in tools.values())} tools")
    
    return registry

from abc import ABC, abstractmethod
from dataclasses import dataclass
from datetime import datetime
from typing import List, AsyncIterator, Optional
from enum import Enum

class SyncMode(Enum):
    """同步模式"""
    FULL = "full"  # 全量同步
    INCREMENTAL = "incremental"  # 增量同步

@dataclass
class DataSourceConfig:
    """数据源配置"""
    source_id: str
    source_type: str  # s3, google_drive, database, etc.
    connection_params: dict
    sync_mode: SyncMode
    sync_schedule: str  # cron 表达式
    filters: dict = None  # 过滤条件
    transformations: List[dict] = None  # 转换规则

@dataclass
class Document:
    """文档对象"""
    doc_id: str
    content: str
    metadata: dict
    source_type: str
    source_id: str
    created_at: datetime
    updated_at: datetime
    checksum: str  # 用于增量检测

class BaseDataConnector(ABC):
    """数据连接器基类"""
    
    def __init__(self, config: DataSourceConfig):
        self.config = config
        self.last_sync_time: Optional[datetime] = None
    
    @abstractmethod
    async def connect(self) -> bool:
        """连接数据源"""
        pass
    
    @abstractmethod
    async def list_documents(
        self,
        since: Optional[datetime] = None
    ) -> List[dict]:
        """列出文档（元数据）"""
        pass
    
    @abstractmethod
    async def fetch_document(self, doc_id: str) -> Document:
        """获取单个文档内容"""
        pass
    
    @abstractmethod
    async def stream_documents(
        self,
        batch_size: int = 100
    ) -> AsyncIterator[List[Document]]:
        """流式获取文档"""
        pass
    
    async def get_changes(self, since: datetime) -> List[dict]:
        """获取变更（用于增量同步）"""
        # 默认实现：比较时间戳
        all_docs = await self.list_documents(since=since)
        return [
            doc for doc in all_docs 
            if doc['updated_at'] > since
        ]

# ============= S3 连接器 =============
class S3Connector(BaseDataConnector):
    """S3 数据连接器"""
    
    async def connect(self) -> bool:
        """连接 S3"""
        self.s3_client = boto3.client('s3', **self.config.connection_params)
        
        # 测试连接
        try:
            bucket = self.config.connection_params['bucket']
            await asyncio.to_thread(
                self.s3_client.head_bucket,
                Bucket=bucket
            )
            return True
        except Exception as e:
            logger.error(f"Failed to connect to S3: {e}")
            return False
    
    async def list_documents(
        self,
        since: Optional[datetime] = None
    ) -> List[dict]:
        """列出 S3 对象"""
        bucket = self.config.connection_params['bucket']
        prefix = self.config.connection_params.get('prefix', '')
        
        paginator = self.s3_client.get_paginator('list_objects_v2')
        
        documents = []
        async for page in self._paginate(paginator, Bucket=bucket, Prefix=prefix):
            for obj in page.get('Contents', []):
                # 过滤时间
                if since and obj['LastModified'] <= since:
                    continue
                
                # 应用过滤器
                if not self._apply_filters(obj):
                    continue
                
                documents.append({
                    'doc_id': obj['Key'],
                    'size': obj['Size'],
                    'updated_at': obj['LastModified'],
                    'etag': obj['ETag']
                })
        
        return documents
    
    async def fetch_document(self, doc_id: str) -> Document:
        """获取 S3 文档内容"""
        bucket = self.config.connection_params['bucket']
        
        # 获取对象元数据
        metadata_response = await asyncio.to_thread(
            self.s3_client.head_object,
            Bucket=bucket,
            Key=doc_id
        )
        
        # 获取对象内容
        content_response = await asyncio.to_thread(
            self.s3_client.get_object,
            Bucket=bucket,
            Key=doc_id
        )
        
        content = content_response['Body'].read()
        
        # 根据文件类型解析
        parsed_content = await self._parse_content(doc_id, content)
        
        return Document(
            doc_id=doc_id,
            content=parsed_content,
            metadata={
                'bucket': bucket,
                'key': doc_id,
                'size': metadata_response['ContentLength'],
                'content_type': metadata_response['ContentType'],
                'etag': metadata_response['ETag']
            },
            source_type='s3',
            source_id=self.config.source_id,
            created_at=metadata_response.get('LastModified'),
            updated_at=metadata_response['LastModified'],
            checksum=metadata_response['ETag']
        )
    
    async def stream_documents(
        self,
        batch_size: int = 100
    ) -> AsyncIterator[List[Document]]:
        """流式获取文档"""
        docs_metadata = await self.list_documents()
        
        batch = []
        for doc_meta in docs_metadata:
            try:
                doc = await self.fetch_document(doc_meta['doc_id'])
                batch.append(doc)
                
                if len(batch) >= batch_size:
                    yield batch
                    batch = []
            except Exception as e:
                logger.error(f"Failed to fetch {doc_meta['doc_id']}: {e}")
                continue
        
        if batch:
            yield batch
    
    def _apply_filters(self, obj: dict) -> bool:
        """应用过滤规则"""
        if not self.config.filters:
            return True
        
        # 文件扩展名过滤
        if 'extensions' in self.config.filters:
            ext = obj['Key'].split('.')[-1].lower()
            if ext not in self.config.filters['extensions']:
                return False
        
        # 大小过滤
        if 'max_size' in self.config.filters:
            if obj['Size'] > self.config.filters['max_size']:
                return False
        
        return True
    
    async def _parse_content(self, key: str, content: bytes) -> str:
        """解析文件内容"""
        ext = key.split('.')[-1].lower()
        
        if ext in ['txt', 'md', 'json', 'csv']:
            return content.decode('utf-8')
        elif ext == 'pdf':
            return await self._parse_pdf(content)
        elif ext in ['docx', 'doc']:
            return await self._parse_word(content)
        elif ext in ['xlsx', 'xls']:
            return await self._parse_excel(content)
        else:
            # 尝试作为文本解析
            try:
                return content.decode('utf-8')
            except:
                return f"[Binary content: {len(content)} bytes]"

# ============= Google Drive 连接器 =============
class GoogleDriveConnector(BaseDataConnector):
    """Google Drive 连接器"""
    
    async def connect(self) -> bool:
        """连接 Google Drive"""
        from google.oauth2.credentials import Credentials
        from googleapiclient.discovery import build
        
        try:
            creds = Credentials.from_authorized_user_info(
                self.config.connection_params['credentials']
            )
            
            self.drive_service = build('drive', 'v3', credentials=creds)
            return True
        except Exception as e:
            logger.error(f"Failed to connect to Google Drive: {e}")
            return False
    
    async def list_documents(
        self,
        since: Optional[datetime] = None
    ) -> List[dict]:
        """列出 Google Drive 文件"""
        query_parts = []
        
        # 文件夹过滤
        if 'folder_id' in self.config.connection_params:
            folder_id = self.config.connection_params['folder_id']
            query_parts.append(f"'{folder_id}' in parents")
        
        # 时间过滤
        if since:
            query_parts.append(
                f"modifiedTime > '{since.isoformat()}'"
            )
        
        # 排除垃圾桶
        query_parts.append("trashed = false")
        
        query = " and ".join(query_parts)
        
        documents = []
        page_token = None
        
        while True:
            results = await asyncio.to_thread(
                self.drive_service.files().list,
                q=query,
                pageSize=1000,
                fields="nextPageToken, files(id, name, mimeType, modifiedTime, size, md5Checksum)",
                pageToken=page_token
            ).execute()
            
            for file in results.get('files', []):
                documents.append({
                    'doc_id': file['id'],
                    'name': file['name'],
                    'mime_type': file['mimeType'],
                    'updated_at': datetime.fromisoformat(
                        file['modifiedTime'].replace('Z', '+00:00')
                    ),
                    'checksum': file.get('md5Checksum')
                })
            
            page_token = results.get('nextPageToken')
            if not page_token:
                break
        
        return documents
    
    async def fetch_document(self, doc_id: str) -> Document:
        """获取 Google Drive 文档"""
        # 获取文件元数据
        file_metadata = await asyncio.to_thread(
            self.drive_service.files().get,
            fileId=doc_id,
            fields='*'
        ).execute()
        
        # 获取文件内容
        mime_type = file_metadata['mimeType']
        
        if 'google-apps' in mime_type:
            # Google 原生文档，需要导出
            content = await self._export_google_doc(doc_id, mime_type)
        else:
            # 普通文件，直接下载
            content = await asyncio.to_thread(
                self.drive_service.files().get_media,
                fileId=doc_id
            ).execute()
            content = content.decode('utf-8')
        
        return Document(
            doc_id=doc_id,
            content=content,
            metadata=file_metadata,
            source_type='google_drive',
            source_id=self.config.source_id,
            created_at=datetime.fromisoformat(
                file_metadata['createdTime'].replace('Z', '+00:00')
            ),
            updated_at=datetime.fromisoformat(
                file_metadata['modifiedTime'].replace('Z', '+00:00')
            ),
            checksum=file_metadata.get('md5Checksum', '')
        )
    
    async def _export_google_doc(self, doc_id: str, mime_type: str) -> str:
        """导出 Google 文档"""
        export_formats = {
            'application/vnd.google-apps.document': 'text/plain',
            'application/vnd.google-apps.spreadsheet': 'text/csv',
            'application/vnd.google-apps.presentation': 'text/plain'
        }
        
        export_mime = export_formats.get(mime_type, 'text/plain')
        
        content = await asyncio.to_thread(
            self.drive_service.files().export_media,
            fileId=doc_id,
            mimeType=export_mime
        ).execute()
        
        return content.decode('utf-8')