#!/usr/bin/env python3
"""
综合MCP测试脚本
测试LiteLLM与MCP服务器的集成
"""

import asyncio
import aiohttp
import json
import time
from typing import Dict, Any, Optional

class MCPTester:
    def __init__(self, litellm_base_url: str = "http://localhost:12168", master_key: str = "sk-1234567890abcdef"):
        self.litellm_base_url = litellm_base_url
        self.master_key = master_key
        self.session = None
        
    async def __aenter__(self):
        self.session = aiohttp.ClientSession()
        return self
        
    async def __aexit__(self, exc_type, exc_val, exc_tb):
        if self.session:
            await self.session.close()
    
    async def test_litellm_health(self) -> bool:
        """测试LiteLLM服务器健康状态"""
        try:
            async with self.session.get(f"{self.litellm_base_url}/health") as response:
                if response.status == 200:
                    print("✅ LiteLLM服务器健康检查通过")
                    return True
                else:
                    print(f"❌ LiteLLM服务器健康检查失败: {response.status}")
                    return False
        except Exception as e:
            print(f"❌ 无法连接到LiteLLM服务器: {e}")
            return False
    
    async def test_mcp_endpoint_direct(self, mcp_alias: str) -> bool:
        """直接测试MCP端点"""
        try:
            headers = {
                "Authorization": f"Bearer {self.master_key}",
                "Content-Type": "application/json"
            }
            
            async with self.session.get(
                f"{self.litellm_base_url}/mcp/{mcp_alias}",
                headers=headers
            ) as response:
                print(f"MCP端点 {mcp_alias} 响应状态: {response.status}")
                
                if response.status == 200:
                    content_type = response.headers.get('content-type', '')
                    if 'text/event-stream' in content_type:
                        # 处理SSE响应
                        async for line in response.content:
                            line_str = line.decode('utf-8').strip()
                            if line_str.startswith('data: '):
                                data = line_str[6:]  # 移除 'data: ' 前缀
                                try:
                                    parsed_data = json.loads(data)
                                    print(f"✅ MCP {mcp_alias} SSE响应: {json.dumps(parsed_data, indent=2)}")
                                    return True
                                except json.JSONDecodeError:
                                    print(f"⚠️  无法解析SSE数据: {data}")
                    else:
                        text = await response.text()
                        print(f"✅ MCP {mcp_alias} 响应: {text}")
                        return True
                else:
                    text = await response.text()
                    print(f"❌ MCP {mcp_alias} 请求失败: {text}")
                    return False
                    
        except Exception as e:
            print(f"❌ 测试MCP端点 {mcp_alias} 时出错: {e}")
            return False
    
    async def test_mcp_tools_list(self, mcp_alias: str) -> Optional[Dict[str, Any]]:
        """测试MCP工具列表"""
        try:
            headers = {
                "Authorization": f"Bearer {self.master_key}",
                "Content-Type": "application/json"
            }
            
            # 构造JSON-RPC请求
            jsonrpc_request = {
                "jsonrpc": "2.0",
                "method": "tools/list",
                "params": {},
                "id": 1
            }
            
            async with self.session.post(
                f"{self.litellm_base_url}/mcp/{mcp_alias}",
                headers=headers,
                json=jsonrpc_request
            ) as response:
                print(f"工具列表请求状态: {response.status}")
                
                if response.status == 200:
                    result = await response.json()
                    print(f"✅ MCP {mcp_alias} 工具列表: {json.dumps(result, indent=2)}")
                    return result
                else:
                    text = await response.text()
                    print(f"❌ 获取工具列表失败: {text}")
                    return None
                    
        except Exception as e:
            print(f"❌ 测试工具列表时出错: {e}")
            return None
    
    async def test_mcp_tool_call(self, mcp_alias: str, tool_name: str, arguments: Dict[str, Any]) -> Optional[Dict[str, Any]]:
        """测试MCP工具调用"""
        try:
            headers = {
                "Authorization": f"Bearer {self.master_key}",
                "Content-Type": "application/json"
            }
            
            # 构造JSON-RPC请求
            jsonrpc_request = {
                "jsonrpc": "2.0",
                "method": "tools/call",
                "params": {
                    "name": tool_name,
                    "arguments": arguments
                },
                "id": 2
            }
            
            async with self.session.post(
                f"{self.litellm_base_url}/mcp/{mcp_alias}",
                headers=headers,
                json=jsonrpc_request
            ) as response:
                print(f"工具调用请求状态: {response.status}")
                
                if response.status == 200:
                    result = await response.json()
                    print(f"✅ MCP {mcp_alias} 工具调用结果: {json.dumps(result, indent=2)}")
                    return result
                else:
                    text = await response.text()
                    print(f"❌ 工具调用失败: {text}")
                    return None
                    
        except Exception as e:
            print(f"❌ 测试工具调用时出错: {e}")
            return None
    
    async def test_direct_mcp_server(self, url: str) -> bool:
        """直接测试MCP服务器"""
        try:
            print(f"\n🔍 直接测试MCP服务器: {url}")
            
            # 测试初始化
            async with self.session.get(url) as response:
                print(f"直接MCP服务器响应状态: {response.status}")
                
                if response.status == 200:
                    content_type = response.headers.get('content-type', '')
                    if 'text/event-stream' in content_type:
                        async for line in response.content:
                            line_str = line.decode('utf-8').strip()
                            if line_str.startswith('data: '):
                                data = line_str[6:]
                                try:
                                    parsed_data = json.loads(data)
                                    print(f"✅ 直接MCP服务器SSE响应: {json.dumps(parsed_data, indent=2)}")
                                    return True
                                except json.JSONDecodeError:
                                    print(f"⚠️  无法解析SSE数据: {data}")
                                break
                    else:
                        text = await response.text()
                        print(f"✅ 直接MCP服务器响应: {text}")
                        return True
                else:
                    text = await response.text()
                    print(f"❌ 直接MCP服务器请求失败: {text}")
                    return False
                    
        except Exception as e:
            print(f"❌ 直接测试MCP服务器时出错: {e}")
            return False
    
    async def run_comprehensive_test(self):
        """运行综合测试"""
        print("🚀 开始MCP综合测试\n")
        
        # 1. 测试LiteLLM健康状态
        print("1️⃣ 测试LiteLLM服务器健康状态")
        health_ok = await self.test_litellm_health()
        
        if not health_ok:
            print("❌ LiteLLM服务器不可用，停止测试")
            return
        
        # 2. 测试本地MCP服务器
        print("\n2️⃣ 测试本地MCP服务器")
        await self.test_direct_mcp_server("http://localhost:8080/mcp")
        
        # 3. 测试通过LiteLLM访问本地MCP
        print("\n3️⃣ 测试通过LiteLLM访问本地MCP")
        test_endpoint_ok = await self.test_mcp_endpoint_direct("test")
        
        if test_endpoint_ok:
            # 4. 测试工具列表
            print("\n4️⃣ 测试本地MCP工具列表")
            tools_result = await self.test_mcp_tools_list("test")
            
            if tools_result and 'result' in tools_result and 'tools' in tools_result['result']:
                tools = tools_result['result']['tools']
                print(f"发现 {len(tools)} 个工具")
                
                # 5. 测试工具调用
                print("\n5️⃣ 测试工具调用")
                for tool in tools[:3]:  # 测试前3个工具
                    tool_name = tool['name']
                    print(f"\n测试工具: {tool_name}")
                    
                    if tool_name == "echo":
                        await self.test_mcp_tool_call("test", "echo", {"message": "Hello MCP!"})
                    elif tool_name == "get_time":
                        await self.test_mcp_tool_call("test", "get_time", {})
                    elif tool_name == "calculate":
                        await self.test_mcp_tool_call("test", "calculate", {"expression": "2+2*3"})
        
        # 6. 测试DeepWiki MCP
        print("\n6️⃣ 测试DeepWiki MCP")
        await self.test_mcp_endpoint_direct("deepwiki")
        
        print("\n🎉 MCP综合测试完成")

async def main():
    """主函数"""
    async with MCPTester() as tester:
        await tester.run_comprehensive_test()

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