#!/usr/bin/env python3
"""
测试MCP服务器的脚本
"""

import asyncio
import json
import subprocess
import sys
from typing import Dict, Any, List

async def test_mcp_server(server_name: str, command: List[str], env: Dict[str, str] = None):
    """测试MCP服务器"""
    print(f"\n=== 测试 {server_name} 服务器 ===")
    
    # 设置环境变量
    process_env = {}
    if env:
        process_env.update(env)
    
    try:
        # 启动服务器进程
        process = await asyncio.create_subprocess_exec(
            *command,
            stdin=asyncio.subprocess.PIPE,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
            env=process_env
        )
        
        # 初始化请求
        init_request = {
            "jsonrpc": "2.0",
            "id": 1,
            "method": "initialize",
            "params": {
                "protocolVersion": "2024-11-05",
                "capabilities": {
                    "tools": {}
                }
            }
        }
        
        # 发送初始化请求
        process.stdin.write((json.dumps(init_request) + "\n").encode())
        await process.stdin.drain()
        
        # 读取初始化响应
        init_response = await process.stdout.readline()
        if init_response:
            try:
                init_data = json.loads(init_response.decode())
                print(f"初始化响应: {init_data}")
            except json.JSONDecodeError:
                print(f"初始化响应解析失败: {init_response}")
        
        # 获取工具列表
        tools_request = {
            "jsonrpc": "2.0",
            "id": 2,
            "method": "tools/list"
        }
        
        # 发送工具列表请求
        process.stdin.write((json.dumps(tools_request) + "\n").encode())
        await process.stdin.drain()
        
        # 读取工具列表响应
        tools_response = await process.stdout.readline()
        if tools_response:
            try:
                tools_data = json.loads(tools_response.decode())
                print(f"工具列表: {json.dumps(tools_data, indent=2, ensure_ascii=False)}")
                
                # 如果有搜索工具，测试搜索功能
                if "result" in tools_data and "tools" in tools_data["result"]:
                    for tool in tools_data["result"]["tools"]:
                        tool_name = tool.get("name")
                        if tool_name and ("search" in tool_name or "document" in tool_name):
                            print(f"\n测试工具: {tool_name}")
                            
                            # 测试搜索工具
                            search_request = {
                                "jsonrpc": "2.0",
                                "id": 3,
                                "method": "tools/call",
                                "params": {
                                    "name": tool_name,
                                    "arguments": {
                                        "query": "测试查询",
                                        "limit": 3
                                    }
                                }
                            }
                            
                            # 发送搜索请求
                            process.stdin.write((json.dumps(search_request) + "\n").encode())
                            await process.stdin.drain()
                            
                            # 读取搜索响应
                            search_response = await process.stdout.readline()
                            if search_response:
                                try:
                                    search_data = json.loads(search_response.decode())
                                    print(f"搜索结果: {json.dumps(search_data, indent=2, ensure_ascii=False)}")
                                except json.JSONDecodeError:
                                    print(f"搜索响应解析失败: {search_response}")
                            break
            except json.JSONDecodeError:
                print(f"工具列表响应解析失败: {tools_response}")
        
        # 关闭进程
        process.stdin.close()
        await process.wait()
        
    except Exception as e:
        print(f"测试 {server_name} 服务器时出错: {e}")

async def main():
    """主函数"""
    print("开始测试MCP服务器...")
    
    # 测试context7服务器
    await test_mcp_server(
        "context7",
        ["npx", "-y", "@upstash/context7-mcp"],
        {"DEFAULT_MINIMUM_TOKENS": ""}
    )
    
    # 测试qdrant服务器
    await test_mcp_server(
        "qdrant",
        ["ssh", "ben@dev1", "cd /home/ben/qdrant && source venv/bin/activate && python qdrant_mcp_server.py"],
        {
            "QDRANT_URL": "http://dev1:6333",
            "QDRANT_API_KEY": "313131",
            "COLLECTION_NAME": "mcp",
            "EMBEDDING_MODEL": "bge-m3"
        }
    )
    
    # 测试qdrant-ollama服务器
    await test_mcp_server(
        "qdrant-ollama",
        ["ssh", "ben@dev1", "cd /home/ben/qdrant && source venv/bin/activate && ./start_mcp_server.sh"],
        {
            "QDRANT_URL": "http://dev1:6333",
            "QDRANT_API_KEY": "313131",
            "COLLECTION_NAME": "ollama_mcp",
            "OLLAMA_MODEL": "nomic-embed-text",
            "OLLAMA_URL": "http://dev1:11434"
        }
    )
    
    print("\n所有测试完成。")

if __name__ == "__main__":
    asyncio.run(main())