package xyz.wbsite.mcp.basic;

import jakarta.annotation.Resource;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.http.MediaType;
import org.springframework.http.ResponseEntity;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestBody;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;

import java.util.Collections;
import java.util.List;
import java.util.Map;

/**
 * 后控制器类
 * 用于处理MCP POST请求
 *
 * @author wangbing
 */
@RestController
@RequestMapping("/mcp/post")
public class PostController {

    private static final Logger log = LoggerFactory.getLogger(PostController.class);
    private static final String WEATHER_TOOL_NAME = "getWeatherForecast";
    private static final String FAKE_WEATHER_JSON = "{\"forecast\": \"sunny\"}";

    @Resource
    SseBroadcaster broadcaster;

    @PostMapping(consumes = MediaType.APPLICATION_JSON_VALUE)
    public ResponseEntity<?> handleMcpPostRequest(@RequestBody McpRequest request) {
        log.info("Received MCP POST Request: ID={}, Method={}", request.getId(), request.getMethod());

        McpResponse mcpResponse = processRequest(request);

        if (mcpResponse != null) {
            // Send the response back over the SSE channel
            broadcaster.broadcastResponse(mcpResponse);
        } else {
            // 处理可能不会生成McpResponse对象的通知等情况
            // 在我们的例子中，'通知/初始化'落在这里。
            log.debug("No explicit response object generated for method '{}', assuming notification ack.", request.getMethod());
        }
        // 立即返回HTTP 200 OK或202 Accepted以确认收到POST。
        // 实际结果通过SSE异步发送。
        // 200 OK可能更简单，因为客户希望得到一些响应体，即使是空的。
        // 202 Accepted明确表示处理正在其他地方进行。让我们用200。
        return ResponseEntity.ok().build();
    }

    private McpResponse processRequest(McpRequest request) {
        switch (request.getMethod()) {
            case "initialize":
                log.info("Handling initialize request");
                InitializeResult initResult = new InitializeResult(new ServerCapabilities());
                return new McpResponse(request.getId(), initResult);

            case "notifications/initialized":
                log.info("Received initialized notification");
                // 这是来自客户端的通知。MCP规范称通知
                // 没有回应。所以我们在这里返回null，POST处理程序
                // 将只返回HTTP OK。
                return null;

            case "tools/list":
                log.info("Handling tools/list request");
                ToolSpecificationData weatherTool = new ToolSpecificationData(
                        WEATHER_TOOL_NAME,
                        "Gets the current weather forecast.",
                        new InputSchema(
                                "object",
                                Map.of("location", Map.of(
                                        "type", "string",
                                        "description", "Location to get the weather for")
                                ),
                                List.of("location"),
                                false)
                );
                ListToolsResult listResult = new ListToolsResult(List.of(weatherTool));
                return new McpResponse(request.getId(), listResult);

            case "tools/call":
                log.info("Handling tools/call request");
                if (request.getParams() != null && request.getParams().has("name")) {
                    String toolName = request.getParams().get("name").asText();
                    if (WEATHER_TOOL_NAME.equals(toolName)) {
                        log.info("Executing tool: {}", toolName);
                        TextContentData textContent = new TextContentData(FAKE_WEATHER_JSON);
                        CallToolResult callResult = new CallToolResult(List.of(textContent));
                        return new McpResponse(request.getId(), callResult);
                    } else {
                        log.warn("Unknown tool requested: {}", toolName);
                        return new McpResponse(request.getId(), new McpError(-32601, "Method not found: " + toolName));
                    }
                } else {
                    log.error("Invalid tools/call request: Missing 'name' in params");
                    return new McpResponse(request.getId(), new McpError(-32602, "Invalid params for tools/call"));
                }

            case "ping":
                log.info("Handling ping request");
                return new McpResponse(request.getId(), Collections.emptyMap());

            default:
                log.warn("Unsupported MCP method: {}", request.getMethod());
                return new McpResponse(request.getId(), new McpError(-32601, "Method not found: " + request.getMethod()));
        }
    }
}