import asyncio
import sys
from openai import AsyncOpenAI
from openai.types.chat import ChatCompletionUserMessageParam
from mcp import ClientSession
from mcp.client.sse import sse_client
from litellm.experimental_mcp_client.tools import (
    transform_mcp_tool_to_openai_tool,
    transform_openai_tool_call_request_to_mcp_tool_call_request,
)


async def main():
    print("开始测试本地MCP连接...")
    
    try:
        # Initialize clients
        print("初始化OpenAI客户端...")
        client = AsyncOpenAI(api_key="sk-1234", base_url="http://localhost:4000")
        print("OpenAI客户端初始化完成")

        # Point MCP client to local LiteLLM Proxy
        print("准备连接本地MCP服务器...")
        
        # 添加超时处理
        try:
            async with asyncio.timeout(10):  # 10秒超时
                print("正在建立SSE连接...")
                async with sse_client("http://localhost:4000/mcp/") as (read, write):
                    print("SSE连接建立成功，初始化会话...")
                    async with ClientSession(read, write) as session:
                        print("正在初始化MCP会话...")
                        await session.initialize()
                        print("MCP会话初始化成功！")

                        # 1. List MCP tools on LiteLLM Proxy
                        print("获取MCP工具列表...")
                        mcp_tools = await session.list_tools()
                        print(f"找到 {len(mcp_tools.tools)} 个MCP工具:")
                        for tool in mcp_tools.tools:
                            print(f"  - {tool.name}: {tool.description}")

                        if not mcp_tools.tools:
                            print("没有找到可用的MCP工具")
                            return

                        # Create message
                        messages = [
                            ChatCompletionUserMessageParam(
                                content="列出所有可用的数据库", role="user"
                            )
                        ]

                        # 2. Convert MCP tools to OpenAI tools
                        print("转换MCP工具为OpenAI格式...")
                        openai_tools = [
                            transform_mcp_tool_to_openai_tool(tool) for tool in mcp_tools.tools
                        ]
                        print(f"转换完成，共 {len(openai_tools)} 个工具")

                        # 3. Call LLM with tools
                        print("调用LLM...")
                        response = await client.chat.completions.create(
                            model="gemini/gemini-2.5-flash",
                            messages=messages,
                            tools=openai_tools,
                            tool_choice="auto",
                        )
                        print("LLM响应完成")

                        # 4. Handle tool call
                        if response.choices[0].message.tool_calls:
                            print("LLM请求调用工具...")
                            tool_call = response.choices[0].message.tool_calls[0]
                            print(f"工具调用: {tool_call.function.name}")
                            print(f"参数: {tool_call.function.arguments}")

                            # 5. Convert to MCP format and execute
                            mcp_call = transform_openai_tool_call_request_to_mcp_tool_call_request(
                                openai_tool=tool_call.model_dump()
                            )

                            print(f"执行MCP工具调用: {mcp_call.name}")
                            result = await session.call_tool(
                                name=mcp_call.name, arguments=mcp_call.arguments
                            )

                            print("工具调用结果:")
                            print(result)
                        else:
                            print("LLM没有请求调用工具")
                            print(f"LLM回复: {response.choices[0].message.content}")
                            
        except asyncio.TimeoutError:
            print("连接超时！可能是本地服务器没有启动或MCP功能未配置")
            return
            
    except Exception as e:
        print(f"发生错误: {type(e).__name__}: {e}")
        import traceback
        traceback.print_exc()


if __name__ == "__main__":
    print("启动本地MCP调试测试...")
    asyncio.run(main())
    print("测试完成")