#!/usr/bin/env python3
"""
Workflow to MCP Server
从 businessQueries.json 或 API 动态生成 MCP 工具

支持两种模式:
- local: 从本地 JSON 文件加载
- api: 从远程 API 加载
"""
import json
import os
import logging
import argparse
import anyio

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, sanitize_json_schema
    from .utils.api_client import execute_workflow, get_workflow_by_id, extract_final_output
    from .utils.env_config import get_workflow_id
    from .utils.logger_config import setup_system_logging, get_logger
except ImportError:
    from schema_converter import convert_sql_params_to_input_schema, sanitize_json_schema
    from utils.api_client import execute_workflow, get_workflow_by_id, extract_final_output
    from utils.env_config import get_workflow_id
    from utils.logger_config import setup_system_logging, get_logger

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

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

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


def parse_arguments():
    """解析命令行参数"""
    parser = argparse.ArgumentParser(description="Workflow MCP Server")
    parser.add_argument(
        "--mode",
        type=str,
        choices=["local", "api"],
        default="api",
        help="数据加载模式: local(本地JSON) 或 api(远程API，默认)"
    )
    parser.add_argument(
        "--json-path",
        type=str,
        default=None,
        help="本地 JSON 文件路径 (local 模式)"
    )
    parser.add_argument(
        "--workflow-id",
        type=str,
        default=None,
        help="工作流ID (api 模式，可选，默认从环境变量 workflowId 获取)"
    )
    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__), "businessQueries.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 远程加载器 - 使用 get_workflow_by_id 获取工作流配置"""
    
    def __init__(self, workflow_id: str = None):
        self.workflow_id = workflow_id or get_workflow_id()
        logger.debug(f"ApiLoader 初始化，工作流ID: {self.workflow_id}")
    
    def load(self) -> list[dict]:
        try:
            if not self.workflow_id:
                logger.error("未提供工作流ID，请设置环境变量 workflowId")
                return []
            
            logger.info(f"开始从 API 加载工作流配置，工作流ID: {self.workflow_id}")
            
            # 调用 get_workflow_by_id 获取工作流配置
            response = get_workflow_by_id(self.workflow_id)
            
            logger.debug(f"API 响应原始数据: {json.dumps(response, ensure_ascii=False, indent=2)}")
            
            # 返回格式: {code, data, msg}
            if response.get("code") != 200:
                logger.error(f"获取工作流配置失败: code={response.get('code')}, msg={response.get('msg')}")
                return []
            
            data = response.get("data")
            logger.debug(f"API 响应 data 字段: {json.dumps(data, ensure_ascii=False, indent=2) if data else 'None'}")
            
            # data 可能是单个对象或列表
            if isinstance(data, dict):
                result = [data]
            elif isinstance(data, list):
                result = data
            else:
                logger.warning(f"API 响应 data 字段类型异常: {type(data)}")
                result = []
            
            logger.info(f"从 API 加载工作流配置成功，工作流ID: {self.workflow_id}, 配置数量: {len(result)}")
            return result
            
        except Exception as e:
            logger.error(f"API 请求失败: {e}", exc_info=True)
            return []


def create_loader(mode: str, **kwargs) -> DataLoader:
    """
    创建数据加载器
    
    Args:
        mode: "local" 或 "api"
        kwargs: 
            - json_path: 本地 JSON 路径 (local 模式)
            - workflow_id: 工作流ID (api 模式，可选，默认从环境变量获取)
    """
    if mode == "local":
        return LocalLoader(json_path=kwargs.get("json_path"))
    elif mode == "api":
        return ApiLoader(workflow_id=kwargs.get("workflow_id"))
    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 query in _tools_config:
        name = query.get("name", "")
        description = query.get("description") or query.get("toolPrompt") or query.get("uniqueName", "")
        
        logger.debug(f"处理工具配置: name={name}, description={description[:50] if description else 'None'}...")
        
        # 优先从 sqlParams 转换，因为它包含更详细的文件配置信息
        sql_params = query.get("sqlParams")
        if sql_params:
            try:
                input_schema = convert_sql_params_to_input_schema(sql_params)
                logger.debug(f"从 sqlParams 转换的 inputSchema: {json.dumps(input_schema, ensure_ascii=False)}")
            except Exception as e:
                logger.warning(f"从 sqlParams 转换失败: {e}，尝试使用 inputJsonSchema")
                # 回退到 inputJsonSchema
                input_json_schema = query.get("inputJsonSchema")
                if input_json_schema:
                    try:
                        if isinstance(input_json_schema, str):
                            input_schema = json.loads(input_json_schema)
                        else:
                            input_schema = input_json_schema
                        input_schema = sanitize_json_schema(input_schema)
                        logger.debug(f"使用 inputJsonSchema (已清理): {json.dumps(input_schema, ensure_ascii=False)}")
                    except (json.JSONDecodeError, TypeError) as e2:
                        logger.error(f"解析 inputJsonSchema 也失败: {e2}，使用空 schema")
                        input_schema = {"type": "object", "properties": {}, "required": []}
                else:
                    logger.warning("sqlParams 和 inputJsonSchema 都不可用，使用空 schema")
                    input_schema = {"type": "object", "properties": {}, "required": []}
        else:
            # 如果没有 sqlParams，尝试使用 inputJsonSchema
            input_json_schema = query.get("inputJsonSchema")
            if input_json_schema:
                try:
                    if isinstance(input_json_schema, str):
                        input_schema = json.loads(input_json_schema)
                    else:
                        input_schema = input_json_schema
                    input_schema = sanitize_json_schema(input_schema)
                    logger.debug(f"使用 inputJsonSchema (已清理): {json.dumps(input_schema, ensure_ascii=False)}")
                except (json.JSONDecodeError, TypeError) as e:
                    logger.error(f"解析 inputJsonSchema 失败: {e}，使用空 schema")
                    input_schema = {"type": "object", "properties": {}, "required": []}
            else:
                logger.warning("sqlParams 和 inputJsonSchema 都不存在，使用空 schema")
                input_schema = {"type": "object", "properties": {}, "required": []}
        
        # 添加 workflow_extraContext 字段到 schema，可以接收任何类型（非必填）
        if "properties" not in input_schema:
            input_schema["properties"] = {}
        input_schema["properties"]["workflow_extraContext"] = {
            "description": "工作流额外的上下文参数（如环境变量等），可以是任何类型，非必填"
        }
        
        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 query in _tools_config:
        if query.get("name") == name:
            tool_config = query
            break
    
    if tool_config is None:
        logger.error(f"未找到工具配置: {name}")
        raise ValueError(f"未知工具: {name}")
    
    logger.debug(f"找到工具配置: {json.dumps(tool_config, ensure_ascii=False, indent=2)}")
    
    # 获取工作流ID
    workflow_id = tool_config.get("workflowId") or get_workflow_id()
    logger.info(f"使用工作流ID: {workflow_id}")
    
    # 提取 workflow_extraContext 字段并合并到 inputs
    inputs = arguments or {}
    workflow_extra_context = inputs.pop("workflow_extraContext", None)
    
    # 如果 workflow_extraContext 存在，将其内容合并到 inputs
    if workflow_extra_context is not None:
        if isinstance(workflow_extra_context, dict):
            # 如果 workflow_extraContext 是字典，合并到 inputs
            inputs.update(workflow_extra_context)
            logger.info(f"workflow_extraContext 是字典类型，已合并到 inputs: {json.dumps(workflow_extra_context, ensure_ascii=False)}")
        else:
            # 如果 workflow_extraContext 不是字典，作为 workflow_extraContext 字段保留
            inputs["workflow_extraContext"] = workflow_extra_context
            logger.info(f"workflow_extraContext 是 {type(workflow_extra_context).__name__} 类型，保留为 workflow_extraContext 字段")
    
    # 构建请求数据
    request_data = {
        "workflowId": workflow_id,
        "inputs": inputs
    }
    
    logger.info(f"执行工作流请求数据: {json.dumps(request_data, ensure_ascii=False, indent=2)}")
    
    try:
        # 调用工作流执行API
        result = execute_workflow(request_data)
        logger.info(f"工作流执行成功: {workflow_id}")
        logger.debug(f"工作流执行结果: {json.dumps(result, ensure_ascii=False, indent=2)}")
        
        # 提取 final_output 最后一个节点的值
        result = extract_final_output(result)
    except Exception as e:
        logger.error(f"工作流执行失败: {e}", exc_info=True)
        result = {
            "error": str(e),
            "tool_name": name,
            "workflow_id": workflow_id
        }
    
    return [
        types.TextContent(
            type="text",
            text=json.dumps(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="workflow_mcp_server",
                server_version="0.1.0",
                capabilities=server.get_capabilities(
                    notification_options=NotificationOptions(),
                    experimental_capabilities={},
                ),
            ),
        )


def main():
    """主入口"""
    global _config
    
    logger.info("=" * 50)
    logger.info("Workflow MCP 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,
        workflow_id=args.workflow_id
    )
    init_tools(loader)
    
    logger.info("开始运行 MCP Server (stdio 模式)")
    # 运行服务器
    anyio.run(run_server)


if __name__ == "__main__":
    main()
