import asyncio
import json
import os
import logging
import argparse
from abc import ABC

import mcp.server.stdio
import mcp.types as types
import requests
from mcp.server import NotificationOptions, Server
from mcp.server.models import InitializationOptions
from omegaconf import OmegaConf

# from src.workflow.workflow_server import WorkflowDifyAPI
from src.difyTaskCall.task_instance import TaskInstance
from src.utils.dify_workflow_schema import process_user_input_form, extract_file_fields
from src.create_mcp_utils import process_file_arguments
from src.utils.logger_config import get_logger
from src.workflow.workflow_server import DifyAPIError

# 使用统一的日志配置
logger = get_logger(__name__)


def parse_arguments():
    parser = argparse.ArgumentParser(description="Dify MCP服务器配置")
    parser.add_argument(
        "--base-url",
        type=str,
        help="API基础URL",
        default="http://192.168.2.236:3001/v1",
    )
    parser.add_argument(
        "--app-sks",
        nargs="+",
        help="应用秘钥列表",
        default=["app-RBS0TuYEnqm8Q1cRQingkuhf"],
    )
    parser.add_argument(
        "--mode-type",
        type=str,
        help="Dify应用模式类型 (workflow, chat, completion)",
        default="workflow",
        choices=["workflow", "chat", "completion"],
    )
    return parser.parse_args()


def get_app_info(base_url=None, app_sks=None, mode_type=None):
    # 获取命令行参数
    args = parse_arguments()
    # 命令行参数优先，其次是函数参数，最后是默认值
    if args.base_url is not None:
        base_url = args.base_url
    if base_url is None:
        base_url = "http://192.168.2.236:3001/v1"

    if args.app_sks is not None:
        app_sks = args.app_sks
    if app_sks is None:
        app_sks = ["app-RBS0TuYEnqm8Q1cRQingkuhf"]

    if args.mode_type is not None:
        mode_type = args.mode_type
    if mode_type is None:
        mode_type = "workflow"

    return base_url, app_sks, mode_type


# 初始化服务器和Dify API
base_url, dify_app_sks, dify_app_mode_type = get_app_info()
server = Server("dify_mcp_server")
task_instance = TaskInstance(base_url, dify_app_sks, dify_app_mode_type)
dify_api = task_instance.get_task_instance(dify_app_mode_type)

# 创建工具
file_config = {
    "file_fields": {},  # 字典，key为工具名称，value为该工具的文件字段列表
    "file_type_dicts": {}  # 字典，key为工具名称，value为该工具的文件类型字典
}


@server.list_tools()
async def handle_list_tools() -> list[types.Tool]:
    """
    列出可用的工具
    返回:
        工具列表，每个工具都使用JSON Schema验证其参数
    """
    tools = []
    tool_names = dify_api.dify_app_names
    tool_infos = dify_api.dify_app_infos
    tool_params = dify_api.dify_app_params
    tool_num = len(tool_names)
    for i in range(tool_num):
        # 加载每个工具的应用信息
        app_info = tool_infos[i]
        # 加载每个工具的应用参数
        app_param = tool_params[i]
        
        # 记录 parameters API 返回的原始数据
        logger.info(f"工具 {app_info['name']} 的 parameters 数据: {app_param}")
        
        # 处理用户输入表单
        inputSchema = process_user_input_form(app_param["user_input_form"])
        # 提取所有文件字段并存储到全局字典中
        tool_file_fields = extract_file_fields(app_param["user_input_form"])
        logger.info(f"工具 {app_info['name']} 提取的文件字段: {tool_file_fields}")
        file_config["file_fields"][app_info["name"]] = tool_file_fields

        

        tools.append(
            types.Tool(
                name=app_info["name"],
                description=app_info["description"],
                inputSchema=inputSchema,
            )
        )
    return tools


@server.call_tool()
async def handle_call_tool(
    name: str, arguments: dict | None
) -> list[types.TextContent | types.ImageContent | types.EmbeddedResource]:
    """
    调用工具处理请求
    参数:
        name: 工具名称
        arguments: 工具参数
    返回:
        处理结果列表
    """
    tool_names = dify_api.dify_app_names
    if name in tool_names:
        tool_idx = tool_names.index(name)
        tool_sk = dify_api.dify_app_sks[tool_idx]

        # 获取当前工具的文件字段信息
        current_tool_file_fields = file_config["file_fields"].get(name, [])
        logger.info(f"工具 {name} 的文件字段信息: {current_tool_file_fields}")
        logger.info(f"工具 {name} 调用前的 arguments: {arguments}")

        # 使用工具函数处理文件字段
        processed_arguments = process_file_arguments(arguments, current_tool_file_fields, dify_api, name)
        logger.info(f"工具 {name} 处理后的 arguments: {processed_arguments}")

        try:
            responses = dify_api.chat_message(
                tool_sk,
                inputs=processed_arguments,
            )

            # 初始化 outputs 变量，避免未定义错误
            outputs = {}
            for res in responses:
                if res["event"] == "workflow_finished":
                    outputs = res["data"]["outputs"]
                    break  # 找到 workflow_finished 事件后退出循环
            
            # 构建 MCP 输出
            mcp_out = []
            if outputs:
                for _, v in outputs.items():
                    mcp_out.append(types.TextContent(type="text", text=str(v)))
            else:
                # 如果没有获取到 outputs，返回错误信息
                logger.warning(f"工具 {name} 未获取到 workflow_finished 事件或 outputs 为空")
                mcp_out.append(types.TextContent(type="text", text="工具执行完成，但未返回输出结果"))
            
            return mcp_out
            
        except DifyAPIError as e:
            # 捕获 Dify API 错误，直接返回给用户
            logger.error(f"工具 {name} 调用 Dify API 失败: {e}")
            error_message = f"API调用失败: {e.message}"
            return [types.TextContent(type="text", text=error_message)]
            
        except requests.exceptions.RequestException as e:
            # 捕获网络请求错误
            logger.error(f"工具 {name} 网络请求失败: {e}")
            error_message = f"网络请求失败: {str(e)}"
            return [types.TextContent(type="text", text=error_message)]
            
        except Exception as e:
            # 捕获其他未知错误
            logger.error(f"工具 {name} 执行时发生未知错误: {e}", exc_info=True)
            error_message = f"工具执行失败: {str(e)}"
            return [types.TextContent(type="text", text=error_message)]
    else:
        raise ValueError(f"Unknown tool: {name}")


def run_main(transport="stdio"):
    """
    主函数：使用stdin/stdout流运行服务器
    """
    if transport == "stdio":
        import anyio
        from mcp.server.stdio import stdio_server

        async def arun():
            async with stdio_server() as streams:
                await server.run(
                    streams[0],
                    streams[1],
                    InitializationOptions(
                        server_name="dify_mcp_server",
                        server_version="0.0.6",
                        capabilities=server.get_capabilities(
                            notification_options=NotificationOptions(),
                            experimental_capabilities={},
                        ),
                    ),
                )

        anyio.run(arun)

    else:
        from mcp.server.sse import SseServerTransport
        from starlette.applications import Starlette
        from starlette.responses import Response
        from starlette.routing import Mount, Route

        sse = SseServerTransport("/messages/")

        async def handle_sse(request):
            async with sse.connect_sse(
                request.scope, request.receive, request._send
            ) as streams:
                await server.run(
                    streams[0], streams[1], server.create_initialization_options()
                )
            return Response()

        starlette_app = Starlette(
            debug=True,
            routes=[
                Route("/sse", endpoint=handle_sse, methods=["GET"]),
                Mount("/messages/", app=sse.handle_post_message),
            ],
        )

        import uvicorn

        uvicorn.run(starlette_app, host="0.0.0.0", port=8080, log_level="info")


if __name__ == "__main__":
    run_main()
