from mcp.server.models import InitializationOptions
from mcp.server import NotificationOptions, Server
from mcp.server.stdio import stdio_server
from mcp.types import (
    Tool,
    TextContent
)
import logging
from src.unitree_adapter import UnitreeRobotAdapter
from src.dji_adapter import DJIDroneAdapter
import json

# 初始化日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

class RobotMCPServer:
    def __init__(self):
        self.unitree = UnitreeRobotAdapter()
        self.dji = DJIDroneAdapter()
        self.server = Server("robot-mcp-server")

    def get_tools(self):
        return [
            Tool(
                name="unitree_connect",
                description="连接宇树机器人",
                inputSchema={
                    "type": "object",
                    "properties": {},
                    "required": []
                }
            ),
            Tool(
                name="unitree_move",
                description="控制宇树机器人运动",
                inputSchema={
                    "type": "object",
                    "properties": {
                        "velocity": {"type": "number", "minimum": 0, "maximum": 3},
                        "duration": {"type": "number", "minimum": 0.1, "default": 1.0}
                    },
                    "required": ["velocity"]
                }
            ),
            Tool(
                name="unitree_stop",
                description="紧急停止宇树机器人",
                inputSchema={
                    "type": "object",
                    "properties": {},
                    "required": []
                }
            ),
            Tool(
                name="dji_connect",
                description="连接大疆无人机",
                inputSchema={
                    "type": "object",
                    "properties": {},
                    "required": []
                }
            ),
            Tool(
                name="dji_takeoff",
                description="大疆无人机起飞",
                inputSchema={
                    "type": "object",
                    "properties": {
                        "height": {"type": "number", "minimum": 1, "maximum": 10, "default": 1.0}
                    },
                    "required": []
                }
            ),
            Tool(
                name="dji_land",
                description="大疆无人机降落",
                inputSchema={
                    "type": "object",
                    "properties": {},
                    "required": []
                }
            )
        ]

    async def handle_tool(self, name: str, arguments: dict):
        try:
            logger.info(f"执行工具 {name}，参数: {arguments}")
            if name == "unitree_connect":
                self.unitree.connect()
                return [TextContent(type="text", text=json.dumps({"status": "connected"}))]
            elif name == "unitree_move":
                velocity = arguments.get("velocity")
                duration = arguments.get("duration", 1.0)
                self.unitree.move(velocity, duration)
                return [TextContent(type="text", text=json.dumps({"velocity": velocity, "duration": duration}))]
            elif name == "unitree_stop":
                self.unitree.emergency_stop()
                return [TextContent(type="text", text=json.dumps({"status": "stopped"}))]
            elif name == "dji_connect":
                self.dji.connect()
                return [TextContent(type="text", text=json.dumps({"status": "connected"}))]
            elif name == "dji_takeoff":
                height = arguments.get("height", 1.0)
                self.dji.takeoff(height)
                return [TextContent(type="text", text=json.dumps({"height": height}))]
            elif name == "dji_land":
                self.dji.land()
                return [TextContent(type="text", text=json.dumps({"status": "landing"}))]
            else:
                raise ValueError(f"未知工具: {name}")
        except Exception as e:
            logger.error(f"工具执行失败: {str(e)}")
            return [TextContent(type="text", text=json.dumps({"error": str(e)}))]

async def main():
    robot_server = RobotMCPServer()
    
    @robot_server.server.list_tools()
    async def list_tools() -> list[Tool]:
        return robot_server.get_tools()
    
    @robot_server.server.call_tool()
    async def call_tool(name: str, arguments: dict) -> list[TextContent]:
        return await robot_server.handle_tool(name, arguments)
    
    async with stdio_server() as (read_stream, write_stream):
        await robot_server.server.run(
            read_stream,
            write_stream,
            InitializationOptions(
                server_name="robot-mcp-server",
                server_version="0.1.0",
                capabilities=robot_server.server.get_capabilities(
                    notification_options=NotificationOptions(),
                    experimental_capabilities={}
                )
            )
        )

def run():
    """入口函数，用于包安装后运行"""
    import asyncio
    asyncio.run(main())

if __name__ == "__main__":
    run()