import argparse
import os
import time

import mcp.types as mt
from fastmcp import FastMCP
from fastmcp.server.middleware import CallNext, Middleware, MiddlewareContext
from main import application
from fastmcp.server.providers.openapi import  MCPType
from fastmcp.tools.base import ToolResult
from config.main import config


def parse_args():
    parser = argparse.ArgumentParser(description='Dcalc MCP server')
    parser.add_argument(
        '--tool-result-mode',
        choices=['structured', 'both'],
        default='structured',
        help='控制工具返回内容：structured 仅保留 structuredContent；both 保留默认双份。',
    )
    args, _ = parser.parse_known_args()
    return args


class ToolResultModeMiddleware(Middleware):
    def __init__(self, mode: str):
        self.mode = mode

    async def on_call_tool(
        self,
        context: MiddlewareContext[mt.CallToolRequestParams],
        call_next: CallNext[mt.CallToolRequestParams, ToolResult],
    ) -> ToolResult:
        result = await call_next(context)
        if result.structured_content is None:
            return result
        if self.mode == 'structured':
            return ToolResult(
                content=[],
                structured_content=result.structured_content,
                meta=result.meta,
            )
        return result


def custom_route_mapper(route, mcp_type: MCPType) -> MCPType | None:
    """Advanced route type mapping."""
    # Convert all admin routes to tools regardless of HTTP method
    if "/open/" in route.path and route.method in ["GET"]:
        return MCPType.TOOL
    else:
        return MCPType.EXCLUDE


def create_mcp(tool_result_mode: str = 'structured') -> FastMCP:
    mcp = FastMCP.from_fastapi(
        app=application,
        name="Dcalc_MCP",
        instructions="DNF相关的MCP服务",
        route_map_fn=custom_route_mapper,
    )
    mcp.add_middleware(ToolResultModeMiddleware(tool_result_mode))
    return mcp


os.environ['TZ'] = 'Asia/Shanghai'
time.tzset()

if __name__ == '__main__':
    args = parse_args()
    mcp = create_mcp(args.tool_result_mode)
    mcp.run(transport="streamable-http", host="0.0.0.0", port=config.MCPPORT)