#!/usr/bin/env python3

import httpx
import json
import asyncio
from typing import AsyncGenerator

async def test_deepwiki_mcp():
    """测试DeepWiki MCP服务器功能"""
    print("=== 测试DeepWiki MCP服务器 ===")
    
    # 测试直接访问DeepWiki MCP端点
    deepwiki_url = "https://mcp.api-inference.modelscope.net/f9d3f201909c45/sse"
    
    try:
        async with httpx.AsyncClient(timeout=30.0) as client:
            print(f"\n1. 测试直接访问DeepWiki MCP端点: {deepwiki_url}")
            
            # 发送SSE请求
            headers = {
                "Accept": "text/event-stream",
                "Cache-Control": "no-cache"
            }
            
            async with client.stream("GET", deepwiki_url, headers=headers) as response:
                print(f"状态码: {response.status_code}")
                print(f"响应头: {dict(response.headers)}")
                
                if response.status_code == 200:
                    print("\n接收到的数据:")
                    count = 0
                    async for line in response.aiter_lines():
                        if line.strip():
                            print(f"Line {count}: {line}")
                            count += 1
                            if count >= 10:  # 限制输出行数
                                print("... (更多数据被截断)")
                                break
                else:
                    print(f"请求失败: {response.status_code}")
                    print(await response.aread())
                    
    except Exception as e:
        print(f"直接访问DeepWiki失败: {e}")
    
    # 测试通过LiteLLM访问DeepWiki MCP
    print("\n\n2. 测试通过LiteLLM访问DeepWiki MCP")
    litellm_mcp_url = "http://localhost:4000/mcp/deepwiki"
    
    try:
        async with httpx.AsyncClient(timeout=30.0) as client:
            headers = {
                "Accept": "text/event-stream",
                "Cache-Control": "no-cache"
            }
            
            async with client.stream("GET", litellm_mcp_url, headers=headers) as response:
                print(f"状态码: {response.status_code}")
                print(f"响应头: {dict(response.headers)}")
                
                if response.status_code == 200:
                    print("\n接收到的数据:")
                    count = 0
                    async for line in response.aiter_lines():
                        if line.strip():
                            print(f"Line {count}: {line}")
                            count += 1
                            if count >= 10:
                                print("... (更多数据被截断)")
                                break
                else:
                    print(f"请求失败: {response.status_code}")
                    error_content = await response.aread()
                    print(f"错误内容: {error_content}")
                    
    except Exception as e:
        print(f"通过LiteLLM访问DeepWiki失败: {e}")
    
    # 测试LiteLLM的基本MCP端点
    print("\n\n3. 测试LiteLLM的基本MCP端点")
    basic_endpoints = [
        "http://localhost:4000/mcp/",
        "http://localhost:4000/mcp",
        "http://localhost:4000/v1/mcp"
    ]
    
    for endpoint in basic_endpoints:
        try:
            async with httpx.AsyncClient(timeout=10.0) as client:
                response = await client.get(endpoint)
                print(f"\n{endpoint}: {response.status_code}")
                if response.status_code != 200:
                    print(f"错误: {response.text[:200]}")
                else:
                    print(f"成功: {response.text[:200]}")
        except Exception as e:
            print(f"\n{endpoint}: 失败 - {e}")

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