#!/usr/bin/env python3
"""
改进的MCP服务器，支持完整的MCP协议
"""

import asyncio
import json
import uuid
from datetime import datetime
from aiohttp import web, web_response
from aiohttp.web import Request, Response

# MCP服务器状态
server_info = {
    "name": "test-mcp-server",
    "version": "1.0.0",
    "protocol_version": "2024-11-05"
}

# 可用工具定义
available_tools = [
    {
        "name": "echo",
        "description": "Echo back the input message",
        "inputSchema": {
            "type": "object",
            "properties": {
                "message": {
                    "type": "string",
                    "description": "Message to echo back"
                }
            },
            "required": ["message"]
        }
    },
    {
        "name": "get_time",
        "description": "Get current time",
        "inputSchema": {
            "type": "object",
            "properties": {},
            "additionalProperties": False
        }
    },
    {
        "name": "calculate",
        "description": "Perform basic arithmetic calculations",
        "inputSchema": {
            "type": "object",
            "properties": {
                "expression": {
                    "type": "string",
                    "description": "Mathematical expression to evaluate (e.g., '2+2', '10*5')"
                }
            },
            "required": ["expression"]
        }
    }
]

async def handle_mcp_request(request: Request) -> Response:
    """处理MCP请求"""
    print(f"收到MCP请求: {request.method} {request.path}")
    print(f"请求头: {dict(request.headers)}")
    
    if request.method == "GET":
        # 处理初始化请求
        return await handle_initialize(request)
    elif request.method == "POST":
        # 处理JSON-RPC请求
        return await handle_jsonrpc(request)
    
    return web_response.Response(status=405, text="Method not allowed")

async def handle_initialize(request: Request) -> Response:
    """处理初始化请求"""
    init_response = {
        "jsonrpc": "2.0",
        "result": {
            "protocolVersion": server_info["protocol_version"],
            "capabilities": {
                "tools": {
                    "listChanged": True
                },
                "resources": {
                    "subscribe": False,
                    "listChanged": False
                },
                "prompts": {
                    "listChanged": False
                },
                "logging": {}
            },
            "serverInfo": {
                "name": server_info["name"],
                "version": server_info["version"]
            }
        },
        "id": 1
    }
    
    # 返回SSE格式的响应
    response_text = f"data: {json.dumps(init_response)}\n\n"
    
    return web_response.Response(
        text=response_text,
        content_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "Connection": "keep-alive",
            "Access-Control-Allow-Origin": "*",
            "Access-Control-Allow-Methods": "GET, POST, OPTIONS",
            "Access-Control-Allow-Headers": "*"
        }
    )

async def handle_jsonrpc(request: Request) -> Response:
    """处理JSON-RPC请求"""
    try:
        body = await request.text()
        print(f"收到JSON-RPC请求体: {body}")
        
        if not body:
            return web_response.Response(status=400, text="Empty request body")
            
        data = json.loads(body)
        method = data.get("method")
        params = data.get("params", {})
        request_id = data.get("id")
        
        print(f"方法: {method}, 参数: {params}")
        
        if method == "tools/list":
            response = {
                "jsonrpc": "2.0",
                "result": {
                    "tools": available_tools
                },
                "id": request_id
            }
        elif method == "tools/call":
            tool_name = params.get("name")
            tool_arguments = params.get("arguments", {})
            
            result = await execute_tool(tool_name, tool_arguments)
            
            response = {
                "jsonrpc": "2.0",
                "result": {
                    "content": [
                        {
                            "type": "text",
                            "text": result
                        }
                    ]
                },
                "id": request_id
            }
        else:
            response = {
                "jsonrpc": "2.0",
                "error": {
                    "code": -32601,
                    "message": f"Method not found: {method}"
                },
                "id": request_id
            }
        
        return web_response.Response(
            text=json.dumps(response),
            content_type="application/json",
            headers={
                "Access-Control-Allow-Origin": "*",
                "Access-Control-Allow-Methods": "GET, POST, OPTIONS",
                "Access-Control-Allow-Headers": "*"
            }
        )
        
    except json.JSONDecodeError as e:
        print(f"JSON解析错误: {e}")
        return web_response.Response(status=400, text="Invalid JSON")
    except Exception as e:
        print(f"处理请求时出错: {e}")
        return web_response.Response(status=500, text="Internal server error")

async def execute_tool(tool_name: str, arguments: dict) -> str:
    """执行工具调用"""
    print(f"执行工具: {tool_name}, 参数: {arguments}")
    
    if tool_name == "echo":
        message = arguments.get("message", "")
        return f"Echo: {message}"
    
    elif tool_name == "get_time":
        current_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        return f"Current time: {current_time}"
    
    elif tool_name == "calculate":
        expression = arguments.get("expression", "")
        try:
            # 简单的数学表达式计算（仅支持基本运算）
            # 注意：这里使用eval有安全风险，实际应用中应该使用更安全的方法
            allowed_chars = set('0123456789+-*/.() ')
            if all(c in allowed_chars for c in expression):
                result = eval(expression)
                return f"Result: {expression} = {result}"
            else:
                return "Error: Invalid characters in expression"
        except Exception as e:
            return f"Error calculating expression: {str(e)}"
    
    else:
        return f"Error: Unknown tool '{tool_name}'"

async def handle_options(request: Request) -> Response:
    """处理OPTIONS请求"""
    return web_response.Response(
        headers={
            "Access-Control-Allow-Origin": "*",
            "Access-Control-Allow-Methods": "GET, POST, OPTIONS",
            "Access-Control-Allow-Headers": "*"
        }
    )

async def create_app():
    """创建web应用"""
    app = web.Application()
    
    # 添加路由
    app.router.add_get('/mcp', handle_mcp_request)
    app.router.add_post('/mcp', handle_mcp_request)
    app.router.add_options('/mcp', handle_options)
    
    return app

if __name__ == '__main__':
    print("启动简单MCP服务器在端口8080...")
    app = asyncio.run(create_app())
    web.run_app(app, host='localhost', port=8080)