#!/usr/bin/env python3
"""
MCP Server 模板
支持从本地 JSON 或 API 动态加载工具配置

使用方式:
- local: 从本地 JSON 文件加载工具配置
- api: 从远程 API 加载工具配置
"""
import json
import os
import logging
import argparse
import anyio
import asyncio

import mcp.types as types
from mcp.server import NotificationOptions, Server
from mcp.server.models import InitializationOptions
from mcp.server.stdio import stdio_server

try:
    from .schema_converter import convert_sql_params_to_input_schema
    from .utils.api_client import call_api
    from .utils.env_config import get_api_key, get_base_url
    from .utils.logger_config import setup_system_logging, get_logger
except ImportError:
    from schema_converter import convert_sql_params_to_input_schema
    from utils.api_client import call_api
    from utils.env_config import get_api_key, get_base_url
    from utils.logger_config import setup_system_logging, get_logger

# 初始化日志系统
setup_system_logging(app_name="lzwcai_mcpskills_template", log_level=logging.DEBUG)
logger = get_logger(__name__)

# 初始化 MCP Server
server = Server("mcpskills_template_server")

# 全局配置
_tools_config: list[dict] = []
_config: dict = {}


def parse_arguments():
    """解析命令行参数"""
    parser = argparse.ArgumentParser(description="MCP Skills Template Server")
    parser.add_argument(
        "--mode",
        type=str,
        choices=["local", "api"],
        default="local",
        help="数据加载模式: local(本地JSON，默认) 或 api(远程API)"
    )
    parser.add_argument(
        "--json-path",
        type=str,
        default=None,
        help="本地 JSON 文件路径 (local 模式)"
    )
    return parser.parse_args()


class DataLoader:
    """数据加载器基类"""
    
    def load(self) -> list[dict]:
        raise NotImplementedError


class LocalLoader(DataLoader):
    """本地 JSON 文件加载器"""
    
    def __init__(self, json_path: str = None):
        if json_path is None:
            json_path = os.path.join(os.path.dirname(__file__), "tools_config.json")
        self.json_path = json_path
    
    def load(self) -> list[dict]:
        try:
            with open(self.json_path, "r", encoding="utf-8") as f:
                data = json.load(f)
                logger.info(f"从本地加载 {len(data)} 条配置: {self.json_path}")
                return data
        except FileNotFoundError:
            logger.error(f"配置文件不存在: {self.json_path}")
            return []
        except json.JSONDecodeError as e:
            logger.error(f"JSON 解析错误: {e}")
            return []


class ApiLoader(DataLoader):
    """API 远程加载器"""
    
    def __init__(self):
        self.base_url = get_base_url()
        self.api_key = get_api_key()
        logger.debug(f"ApiLoader 初始化，base_url: {self.base_url}")
    
    def load(self) -> list[dict]:
        try:
            # TODO: 根据实际 API 修改此处
            logger.info(f"开始从 API 加载工具配置")
            response = call_api("/api/tools", method="GET")
            
            if response.get("code") != 200:
                logger.error(f"获取工具配置失败: {response.get('msg')}")
                return []
            
            data = response.get("data", [])
            if isinstance(data, dict):
                data = [data]
            
            logger.info(f"从 API 加载工具配置成功，数量: {len(data)}")
            return data
            
        except Exception as e:
            logger.error(f"API 请求失败: {e}", exc_info=True)
            return []


def create_loader(mode: str, **kwargs) -> DataLoader:
    """创建数据加载器"""
    if mode == "local":
        return LocalLoader(json_path=kwargs.get("json_path"))
    elif mode == "api":
        return ApiLoader()
    else:
        raise ValueError(f"不支持的模式: {mode}")


def init_tools(loader: DataLoader):
    """初始化工具配置"""
    global _tools_config
    _tools_config = loader.load()
    logger.info(f"已加载 {len(_tools_config)} 个工具配置")


@server.list_tools()
async def handle_list_tools() -> list[types.Tool]:
    """列出所有可用工具"""
    logger.info(f"收到 ListTools 请求，当前配置数量: {len(_tools_config)}")
    tools = []
    
    for tool in _tools_config:
        name = tool.get("name", "")
        description = tool.get("description") or tool.get("toolPrompt", "")
        sql_params = tool.get("sqlParams", "[]")
        
        # 转换参数为 inputSchema
        input_schema = convert_sql_params_to_input_schema(sql_params)
        
        tools.append(
            types.Tool(
                name=name,
                description=description,
                inputSchema=input_schema
            )
        )
    
    logger.info(f"ListTools 响应: 返回 {len(tools)} 个工具")
    return tools


@server.call_tool()
async def handle_call_tool(
    name: str, 
    arguments: dict | None
) -> list[types.TextContent | types.ImageContent | types.EmbeddedResource]:
    """调用工具"""
    logger.info(f"收到 CallTool 请求: name={name}, arguments={json.dumps(arguments, ensure_ascii=False) if arguments else 'None'}")
    
    # 查找对应的工具配置
    tool_config = None
    for tool in _tools_config:
        if tool.get("name") == name:
            tool_config = tool
            break
    
    if tool_config is None:
        logger.error(f"未找到工具配置: {name}")
        raise ValueError(f"未知工具: {name}")
    
    # TODO: 在这里实现你的工具逻辑
    try:
        # execute_tool 现在是 async 的
        result_contents = await execute_tool(name, arguments or {}, tool_config)
        logger.info(f"工具执行成功: {name}")
        return result_contents
    except Exception as e:
        logger.error(f"工具执行失败: {e}", exc_info=True)
        # 重新抛出异常，以便 MCP SDK 能够捕获并将其格式化为错误响应
        raise e


async def execute_tool(name: str, arguments: dict, config: dict) -> list[types.TextContent | types.ImageContent | types.EmbeddedResource]:
    """
    执行工具逻辑并返回 MCP 内容列表 (Async)
    """
    # 1. 示例：Hello World (返回纯文本)
    if name == "example_hello_world":
        msg = f"Hello, {arguments.get('name', 'World')}!"
        return [types.TextContent(type="text", text=msg)]
    
    # 2. 示例：计算器 (返回 JSON 格式的文本)
    elif name == "example_calculator":
        a = float(arguments.get("a", 0))
        b = float(arguments.get("b", 0))
        op = arguments.get("operation", "+")
        
        if op == "+":
            result = a + b
        elif op == "-":
            result = a - b
        elif op == "*":
            result = a * b
        elif op == "/":
            result = a / b if b != 0 else "除数不能为0"
        else:
            result = "未知运算符"
        
        output = {"result": result, "expression": f"{a} {op} {b}"}
        return [types.TextContent(type="text", text=json.dumps(output, ensure_ascii=False))]
    
    # 3. 演示：纯文本
    elif name == "demo_text":
        return [types.TextContent(type="text", text="这是一个标准的纯文本输出示例。")]
        
    # 4. 演示：图片输出
    elif name == "demo_image":
        # 这是一个 1x1 的红色像素点的 Base64
        red_dot_base64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z8BQDwAEhQGAhKmMIQAAAABJRU5ErkJggg=="
        return [
            types.ImageContent(
                type="image",
                data=red_dot_base64,
                mimeType="image/png"
            )
        ]

    # 5. 演示：资源输出
    elif name == "demo_resource":
        # 模拟返回一个嵌入式资源
        return [
            types.EmbeddedResource(
                type="resource",
                resource=types.TextResourceContents(
                    uri="file:///logs/app.log",
                    text="[INFO] System started\n[WARN] Low memory",
                    mimeType="text/plain"
                )
            )
        ]

    # 6. 演示：错误状态 (带 observation)
    elif name == "demo_error":
        # 尝试模拟用户要求的带 observation 的错误结构
        # 注意：标准 MCP 协议中 TextContent 可能不支持 observation 字段
        # 这里演示如何抛出异常，这是最标准的错误反馈方式
        raise ValueError("执行失败: 网络超时，请检查连接")
        
        # 如果客户端支持非标准字段，可以尝试返回如下结构（需 SDK 支持）：
        # return [
        #    types.TextContent(type="text", text="执行失败")  # 无法直接添加 observation
        # ]

    # 7. 演示：混合输出 (文本 + 图片)
    elif name == "demo_mixed":
        red_dot_base64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z8BQDwAEhQGAhKmMIQAAAABJRU5ErkJggg=="
        return [
            types.TextContent(type="text", text="下面是一张图片："),
            types.ImageContent(
                type="image",
                data=red_dot_base64,
                mimeType="image/png"
            ),
            types.TextContent(type="text", text="图片展示完毕。")
        ]
        
    # 8. 演示：复杂参数
    elif name == "demo_complex":
        user_info = arguments.get("user_info", {})
        tags = arguments.get("tags", [])
        
        # 模拟异步处理
        await asyncio.sleep(0.1)
        
        result_text = f"接收到复杂参数:\n用户: {json.dumps(user_info, ensure_ascii=False)}\n标签: {tags}"
        return [types.TextContent(type="text", text=result_text)]
    
    # 默认返回参数
    default_result = {
        "tool_name": name,
        "arguments": arguments,
        "message": "工具执行成功（默认实现）"
    }
    return [types.TextContent(type="text", text=json.dumps(default_result, ensure_ascii=False, indent=2))]


async def run_server():
    """运行 MCP Server (stdio 模式)"""
    async with stdio_server() as streams:
        await server.run(
            streams[0],
            streams[1],
            InitializationOptions(
                server_name="mcpskills_template_server",
                server_version="0.1.0",
                capabilities=server.get_capabilities(
                    notification_options=NotificationOptions(),
                    experimental_capabilities={},
                ),
            ),
        )


def main():
    """主入口"""
    global _config
    
    logger.info("=" * 50)
    logger.info("MCP Skills Template Server 启动")
    logger.info("=" * 50)
    
    # 解析命令行参数
    args = parse_arguments()
    _config = vars(args)
    logger.info(f"命令行参数: {_config}")
    
    # 创建加载器并初始化工具
    logger.info(f"使用模式: {args.mode}")
    loader = create_loader(
        mode=args.mode,
        json_path=args.json_path
    )
    init_tools(loader)
    
    logger.info("开始运行 MCP Server (stdio 模式)")
    # 运行服务器
    anyio.run(run_server)


if __name__ == "__main__":
    main()
