#!/usr/bin/env python3
"""
专门测试qdrant-ollama服务器的脚本
"""

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

def test_qdrant_ollama_server():
    """测试qdrant-ollama服务器"""
    print("\n=== 测试 qdrant-ollama 服务器 ===")
    
    try:
        # 启动服务器进程
        process = subprocess.Popen(
            ["ssh", "ben@dev1", "cd /home/ben/qdrant && source venv/bin/activate && ./start_mcp_server.sh"],
            stdin=subprocess.PIPE,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True
        )
        
        # 读取并忽略所有非JSON输出
        buffer = ""
        json_found = False
        
        # 等待进程启动并读取初始输出
        for _ in range(10):  # 最多尝试10次
            line = process.stdout.readline()
            if not line:
                time.sleep(0.5)
                continue
            
            line = line.strip()
            buffer += line + "\n"
            
            # 尝试解析JSON
            try:
                data = json.loads(line)
                if "jsonrpc" in data:
                    json_found = True
                    print(f"收到JSON响应: {json.dumps(data, indent=2, ensure_ascii=False)}")
                    break
            except json.JSONDecodeError:
                # 不是JSON，继续读取
                continue
        
        if not json_found:
            print(f"未找到JSON响应，原始输出: {buffer}")
            process.terminate()
            process.wait()
            return
        
        # 初始化请求
        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}")
                            
                            # 先添加一个文档
                            add_request = {
                                "jsonrpc": "2.0",
                                "id": 3,
                                "method": "tools/call",
                                "params": {
                                    "name": "add_document",
                                    "arguments": {
                                        "text": "这是一个测试文档，用于验证qdrant-ollama服务器的功能。",
                                        "metadata": {
                                            "source": "test",
                                            "topic": "测试"
                                        }
                                    }
                                }
                            }
                            
                            # 发送添加文档请求
                            process.stdin.write(json.dumps(add_request) + "\n")
                            process.stdin.flush()
                            
                            # 读取添加文档响应
                            add_response = process.stdout.readline()
                            if add_response:
                                try:
                                    add_data = json.loads(add_response.strip())
                                    print(f"添加文档测试成功")
                                except json.JSONDecodeError:
                                    print(f"添加文档响应解析失败: {add_response}")
                            
                            # 测试搜索工具
                            search_request = {
                                "jsonrpc": "2.0",
                                "id": 4,
                                "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"测试 qdrant-ollama 服务器时出错: {e}")

def main():
    """主函数"""
    print("开始测试qdrant-ollama服务器...")
    
    test_qdrant_ollama_server()
    
    print("\n测试完成。")

if __name__ == "__main__":
    main()