import asyncio
import logging
from typing import Dict, List, Optional, Any, Callable, Union
import os

# 导入MCP SDK相关库
from mcp.server.fastmcp import FastMCP
from mcp.types import Resource, Tool, Prompt
from mcp import types

class MCPServer:
    """
    Model Context Protocol 服务器
    用于管理模型连接和资源交互
    """
    instance = None

    def __init__(self, server_name: str = "ComfyUI-MCP"):
        """
        初始化MCP服务器
        
        参数:
            server_name: 服务器名称
        """
        MCPServer.instance = self
        
        # 初始化FastMCP实例
        self.mcp = FastMCP(server_name)
        
        # 存储注册的资源、工具和提示
        self.registered_resources: Dict[str, Callable] = {}
        self.registered_tools: Dict[str, Callable] = {}
        self.registered_prompts: Dict[str, Callable] = {}
        
        # 存储模型连接信息
        self.active_models: Dict[str, Dict[str, Any]] = {}
        self.model_count: int = 0
        
        # 状态标志
        self.is_running: bool = False
        
        # 设置基本资源和工具
        self._setup_default_handlers()
        
        logging.info(f"[MCP服务器] 已初始化，服务器名称: {server_name}")

    def _setup_default_handlers(self):
        """设置默认的资源处理器和工具"""
        
        # 注册状态资源
        @self.mcp.resource("status://server")
        def get_server_status() -> str:
            """返回服务器状态信息"""
            return f"""
            服务器状态:
            - 运行中: {self.is_running}
            - 活跃模型数: {len(self.active_models)}
            - 已注册资源数: {len(self.registered_resources)}
            - 已注册工具数: {len(self.registered_tools)}
            - 已注册提示数: {len(self.registered_prompts)}
            """
        
        # 注册模型列表资源
        @self.mcp.resource("models://list")
        def list_models() -> str:
            """返回当前连接的模型列表"""
            if not self.active_models:
                return "当前没有活跃的模型连接"
            
            result = "活跃模型列表:\n"
            for model_id, model_info in self.active_models.items():
                result += f"- ID: {model_id}, 名称: {model_info.get('name', 'Unknown')}\n"
            return result
        
        # 注册Echo工具
        @self.mcp.tool()
        def echo(message: str) -> str:
            """简单的Echo工具，用于测试连接"""
            return f"MCP服务器回声: {message}"
        
        # 注册系统信息工具
        @self.mcp.tool()
        def system_info() -> dict:
            """返回系统信息"""
            import platform
            return {
                "os": platform.system(),
                "python_version": platform.python_version(),
                "hostname": platform.node(),
                "cpu": platform.processor()
            }

    def register_resource(self, uri_pattern: str):
        """
        注册一个资源处理函数
        
        参数:
            uri_pattern: 资源URI模式
        """
        def decorator(func):
            self.registered_resources[uri_pattern] = func
            self.mcp.resource(uri_pattern)(func)
            logging.info(f"[MCP服务器] 已注册资源: {uri_pattern}")
            return func
        return decorator

    def register_tool(self, name: Optional[str] = None):
        """
        注册一个工具函数
        
        参数:
            name: 工具名称（可选）
        """
        def decorator(func):
            tool_name = name or func.__name__
            self.registered_tools[tool_name] = func
            self.mcp.tool(name=tool_name)(func)
            logging.info(f"[MCP服务器] 已注册工具: {tool_name}")
            return func
        return decorator

    def register_prompt(self, name: Optional[str] = None):
        """
        注册一个提示模板
        
        参数:
            name: 提示名称（可选）
        """
        def decorator(func):
            prompt_name = name or func.__name__
            self.registered_prompts[prompt_name] = func
            self.mcp.prompt(name=prompt_name)(func)
            logging.info(f"[MCP服务器] 已注册提示: {prompt_name}")
            return func
        return decorator

    def register_model(self, model_info: Dict[str, Any]) -> str:
        """
        注册一个模型到MCP服务器
        
        参数:
            model_info: 模型信息字典
            
        返回:
            model_id: 模型ID
        """
        model_id = f"model_{self.model_count}"
        self.model_count += 1
        
        self.active_models[model_id] = {
            "id": model_id,
            "registered_at": asyncio.get_event_loop().time(),
            **model_info
        }
        
        logging.info(f"[MCP服务器] 已注册模型: {model_id}")
        return model_id

    def unregister_model(self, model_id: str) -> bool:
        """
        从MCP服务器注销一个模型
        
        参数:
            model_id: 模型ID
            
        返回:
            成功与否
        """
        if model_id in self.active_models:
            del self.active_models[model_id]
            logging.info(f"[MCP服务器] 已注销模型: {model_id}")
            return True
        
        logging.warning(f"[MCP服务器] 尝试注销未知模型: {model_id}")
        return False

    async def start(self, host: str = "127.0.0.1", port: int = 8189):
        """
        启动MCP服务器
        
        参数:
            host: 主机地址
            port: 端口号
        """
        self.is_running = True
        try:
            # 启动FastMCP HTTP服务器
            uvicorn_config = {
                "host": host,
                "port": port,
                "log_level": "info"
            }
            
            logging.info(f"[MCP服务器] 正在启动... 地址: {host}:{port}")
            
            # 使用FastMCP的HTTP服务器启动方法
            await self.mcp.serve(**uvicorn_config)
        except Exception as e:
            self.is_running = False
            logging.error(f"[MCP服务器] 启动失败: {str(e)}")
            raise
        finally:
            self.is_running = False

    def get_stats(self) -> Dict[str, Any]:
        """获取服务器统计信息"""
        return {
            "is_running": self.is_running,
            "active_models_count": len(self.active_models),
            "resources_count": len(self.registered_resources),
            "tools_count": len(self.registered_tools),
            "prompts_count": len(self.registered_prompts)
        }

# 运行MCP服务器的异步函数
async def run_server(server_name: str = "ComfyUI-MCP", host: str = "127.0.0.1", port: int = 8189):
    """
    运行MCP服务器的便捷函数
    
    参数:
        server_name: 服务器名称
        host: 主机地址
        port: 端口号
    """
    server = MCPServer(server_name)
    await server.start(host, port)
    return server

# 同步运行MCP服务器的函数
def run(server_name: str = "ComfyUI-MCP", host: str = "127.0.0.1", port: int = 8189):
    """
    同步运行MCP服务器
    
    参数:
        server_name: 服务器名称
        host: 主机地址
        port: 端口号
    """
    loop = asyncio.get_event_loop()
    try:
        return loop.run_until_complete(run_server(server_name, host, port))
    except KeyboardInterrupt:
        logging.info("[MCP服务器] 收到中断信号，正在关闭...")
    finally:
        loop.close()