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
# 配置日志记录
logging.basicConfig(
    level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger(__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-XaRWpeL2Yfdguc5ul3ScXvPE"],
    )
    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-XaRWpeL2Yfdguc5ul3ScXvPE"]

    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)


def process_user_input_form1(user_input_form):
    """
    处理Dify应用的用户输入表单，转换为JSON Schema格式

    参数:
        user_input_form: Dify应用的用户输入表单配置

    返回:
        处理后的inputSchema字典
    """
    inputSchema = dict(
        type="object",
        properties={},
        required=[],
    )

    property_num = len(user_input_form)
    if property_num > 0:
        for j in range(property_num):
            param = user_input_form[j]
            param_type = list(param.keys())[0]
            param_info = param[param_type]
            property_name = param_info["variable"]

            # 根据不同控件类型处理
            if param_type == "text-input":
                inputSchema["properties"][property_name] = {
                    "type": "string",
                    "description": param_info["label"],
                }
                if "default" in param_info:
                    inputSchema["properties"][property_name]["default"] = param_info[
                        "default"
                    ]

            elif param_type == "paragraph":
                inputSchema["properties"][property_name] = {
                    "type": "string",
                    "description": param_info["label"],
                    "format": "paragraph",
                }
                if "default" in param_info:
                    inputSchema["properties"][property_name]["default"] = param_info[
                        "default"
                    ]

            elif param_type == "select":
                inputSchema["properties"][property_name] = {
                    "type": "string",
                    "description": param_info["label"],
                    "enum": param_info["options"],
                }
                if "default" in param_info:
                    inputSchema["properties"][property_name]["default"] = param_info[
                        "default"
                    ]

            elif param_type == "file_upload":
                # 文件上传控件处理
                file_type_schema = {
                    "type": "object",
                    "description": param_info["label"],
                    "properties": {
                        "file_url": {"type": "string", "description": "文件URL"},
                        "file_name": {"type": "string", "description": "文件名称"},
                    },
                    "required": ["file_url"],
                }

                # 处理图片上传配置
                if "image" in param_info and param_info["image"]["enabled"]:
                    image_config = param_info["image"]
                    file_type_schema["properties"]["type"] = {
                        "type": "string",
                        "description": "文件类型，支持png、jpg、jpeg、webp、gif",
                        "enum": ["png", "jpg", "jpeg", "webp", "gif"],
                    }

                    # 处理数量限制
                    number_limits = image_config.get("number_limits", 3)
                    if number_limits > 1:
                        # 如果允许多个文件，则使用数组
                        inputSchema["properties"][property_name] = {
                            "type": "array",
                            "description": param_info["label"],
                            "items": file_type_schema,
                            "maxItems": number_limits,
                        }
                    else:
                        # 如果只允许单个文件
                        inputSchema["properties"][property_name] = file_type_schema
                else:
                    # 如果没有特定的图片配置，使用一般文件配置
                    inputSchema["properties"][property_name] = file_type_schema

            else:
                # 默认处理为字符串类型
                inputSchema["properties"][property_name] = {
                    "type": "string",
                    "description": param_info["label"],
                }

            # 处理必填字段
            if param_info.get("required", False):
                inputSchema["required"].append(property_name)

    # 添加必填的userId参数，支持数字或字符串类型
    inputSchema["properties"]["userId"] = dict(
        oneOf=[{"type": "number"}, {"type": "string"}],
        description="您的员工ID，用于识别您的员工身份",
    )
    inputSchema["required"].append("userId")

    return inputSchema


@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]
        # 处理用户输入表单
        inputSchema = process_user_input_form(app_param["user_input_form"])

        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]

        # 提取files参数，并创建不包含files的inputs对象
        files = arguments.get("files", None) if arguments else None
        inputs = {k: v for k, v in (arguments or {}).items() if k != "files"}

        responses = dify_api.chat_message(
            tool_sk,
            inputs=inputs,
            userId=arguments.get("userId", "pp666") if arguments else "pp666",
            files=files,
        )

  
        for res in responses:
            if res["event"] == "workflow_finished":
                outputs = res["data"]["outputs"]
        mcp_out = []
        for _, v in outputs.items():
            mcp_out.append(types.TextContent(type="text", text=v))
        return mcp_out
    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()
