#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
MCP服务器示例脚本
---------------
这个脚本演示了如何启动和使用MCP服务器
"""

import logging
import asyncio
from typing import Dict, Any

# 设置日志级别
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)

# 导入MCP服务器
from mcp_server import MCPServer

# 自定义资源和工具示例
async def setup_custom_handlers(server: MCPServer):
    """设置自定义资源和工具"""
    
    # 注册自定义资源
    @server.register_resource("comfyui://status")
    def get_comfy_status() -> str:
        """获取ComfyUI状态信息"""
        return "ComfyUI 正在运行中"
    
    # 注册自定义工具
    @server.register_tool()
    def generate_image(prompt: str, width: int = 512, height: int = 512) -> Dict[str, Any]:
        """生成图像的工具（示例）"""
        # 这里只是示例，实际实现需要与ComfyUI集成
        return {
            "status": "success",
            "prompt": prompt,
            "dimensions": f"{width}x{height}",
            "message": "图像生成请求已提交"
        }
    
    # 注册自定义提示
    @server.register_prompt("图像生成")
    def image_generation_prompt(style: str, subject: str) -> str:
        """创建图像生成提示"""
        return f"请创建一个{style}风格的{subject}图像。"

async def register_models(server: MCPServer):
    """注册示例模型"""
    # 注册示例模型
    model1 = server.register_model({
        "name": "stable-diffusion-v1.5",
        "type": "diffusion",
        "description": "Stable Diffusion v1.5图像生成模型"
    })
    
    model2 = server.register_model({
        "name": "llama-3-8b",
        "type": "llm",
        "description": "Llama 3 8B语言模型"
    })
    
    logging.info(f"已注册模型: {model1}, {model2}")
    return [model1, model2]

async def main():
    """主函数"""
    # 创建MCP服务器实例
    server = MCPServer(server_name="ComfyUI-MCP示例")
    
    # 设置自定义处理器
    await setup_custom_handlers(server)
    
    # 注册模型
    models = await register_models(server)
    
    # 打印服务器信息
    logging.info(f"服务器统计信息: {server.get_stats()}")
    
    # 启动服务器（默认地址为127.0.0.1:8189）
    try:
        logging.info("正在启动MCP服务器...")
        await server.start()
    except KeyboardInterrupt:
        logging.info("服务器被用户中断")
    finally:
        # 清理资源
        for model_id in models:
            server.unregister_model(model_id)
        logging.info("服务器已关闭")

if __name__ == "__main__":
    # 运行主函数
    asyncio.run(main())