package com.zy.ai.mcp.controller;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONObject;
import com.zy.ai.mcp.config.McpToolsBootstrap;
import com.zy.ai.mcp.dto.JsonRpcRequest;
import com.zy.ai.mcp.dto.JsonRpcResponse;
import com.zy.ai.mcp.dto.ToolDefinition;
import com.zy.ai.mcp.dto.ToolRegistry;
import com.zy.ai.mcp.service.WcsDataFacade;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.MediaType;
import org.springframework.web.bind.annotation.*;

import javax.annotation.PostConstruct;
import java.util.*;

@Slf4j
@RestController
@RequestMapping("/ai/mcp")
public class McpController {

    private final ToolRegistry registry = new ToolRegistry();

    @Autowired
    private WcsDataFacade wcsDataFacade;

    public McpController(WcsDataFacade wcsDataFacade) {
        this.wcsDataFacade = wcsDataFacade;
    }

    @PostConstruct
    public void init() {
        McpToolsBootstrap.registerAll(registry, wcsDataFacade);
        log.info("MCP initialized, tools={}", registry.listTools().size());
    }

    @PostMapping(consumes = MediaType.APPLICATION_JSON_VALUE, produces = MediaType.APPLICATION_JSON_VALUE)
    public Object handle(@RequestBody JsonRpcRequest req,
                         @RequestHeader(value = "Authorization", required = false) String auth) {

        // （建议）做一个简单鉴权：防止被随便调用生产系统
        // if (!"Bearer your-token".equals(auth)) return JsonRpcResponse.err(null, 401, "Unauthorized", null);

        String id = req.getId();
        String method = req.getMethod();
        JSONObject params = JSON.parseObject(JSON.toJSONString(req.getParams()));

        try {
            if ("initialize".equals(method)) {
                Map<String, Object> result = new LinkedHashMap<String, Object>();
                result.put("serverName", "wcs-mcp");
                result.put("serverVersion", "1.0.0");
                result.put("capabilities", Arrays.asList("tools"));
                return JsonRpcResponse.ok(id, result);
            }

            if ("tools/list".equals(method)) {
                Map<String, Object> result = new LinkedHashMap<String, Object>();
                result.put("tools", registry.listTools());
                // cursor/paging 你后面需要再加
                return JsonRpcResponse.ok(id, result);
            }

            if ("tools/call".equals(method)) {
                String toolName = params.getString("name");
                JSONObject arguments = params.getJSONObject("arguments");
                if (toolName == null || toolName.trim().isEmpty()) {
                    return JsonRpcResponse.err(id, -32602, "Invalid params: missing tool name", null);
                }
                ToolDefinition def = registry.get(toolName);
                if (def == null) {
                    return JsonRpcResponse.err(id, -32601, "Method not found: tool " + toolName, null);
                }
                Object output = def.getHandler().handle(arguments == null ? new JSONObject() : arguments);

                Map<String, Object> result = new LinkedHashMap<String, Object>();
                result.put("content", output); // 你也可以按 MCP 常见返回结构做 text/json 分段
                return JsonRpcResponse.ok(id, result);
            }

            return JsonRpcResponse.err(id, -32601, "Method not found: " + method, null);

        } catch (Exception e) {
            log.error("MCP handle error, method={}, params={}", method, params, e);
            return JsonRpcResponse.err(id, -32000, "Server error", e.getMessage());
        }
    }
}