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

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

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 = subprocess.Popen(
            command,
            stdin=subprocess.PIPE,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            env=process_env,
            text=True
        )
        
        # 等待进程启动
        time.sleep(2)
        
        # 初始化请求
        init_request = {
            "jsonrpc": "2.0",
            "id": 1,
            "method": "initialize",
            "params": {
                "protocolVersion": "2024-11-05",
                "capabilities": {
                    "tools": {}
                }
            }
        }
        
        # 发送初始化请求
        process.stdin.write(json.dumps(init_request) + "\n")
        process.stdin.flush()
        
        # 读取初始化响应
        init_response = process.stdout.readline()
        if init_response:
            try:
                init_data = json.loads(init_response.strip())
                print(f"初始化成功: {init_data.get('result', {}).get('serverInfo', {}).get('name', '未知服务器')}")
            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")
        process.stdin.flush()
        
        # 读取工具列表响应
        tools_response = process.stdout.readline()
        if tools_response:
            try:
                tools_data = json.loads(tools_response.strip())
                print(f"工具列表获取成功")
                
                # 如果有搜索工具，测试搜索功能
                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")
                            process.stdin.flush()
                            
                            # 读取搜索响应
                            search_response = process.stdout.readline()
                            if search_response:
                                try:
                                    search_data = json.loads(search_response.strip())
                                    print(f"搜索测试成功")
                                    if "result" in search_data and "content" in search_data["result"]:
                                        for content in search_data["result"]["content"]:
                                            if content.get("type") == "text":
                                                print(f"搜索结果: {content.get('text', '')[:100]}...")
                                except json.JSONDecodeError:
                                    print(f"搜索响应解析失败: {search_response}")
                            break
            except json.JSONDecodeError:
                print(f"工具列表响应解析失败: {tools_response}")
        
        # 关闭进程
        process.stdin.close()
        process.terminate()
        process.wait()
        
    except Exception as e:
        print(f"测试 {server_name} 服务器时出错: {e}")

def main():
    """主函数"""
    print("开始测试MCP服务器...")
    
    # 测试context7服务器
    test_mcp_server(
        "context7",
        ["npx", "-y", "@upstash/context7-mcp"],
        {"DEFAULT_MINIMUM_TOKENS": ""}
    )
    
    # 测试qdrant服务器
    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服务器
    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__":
    main()