package com.vincent.rsf.server.ai.service.mcp;

import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.vincent.rsf.server.ai.constant.AiMcpConstants;
import com.vincent.rsf.server.ai.model.AiDiagnosticToolResult;
import com.vincent.rsf.server.ai.model.AiMcpToolDescriptor;
import com.vincent.rsf.server.ai.model.AiPromptContext;
import org.springframework.stereotype.Service;

import javax.annotation.Resource;
import java.util.ArrayList;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;

@Service
public class AiMcpProtocolService {

    @Resource
    private AiMcpRegistryService aiMcpRegistryService;
    @Resource
    private ObjectMapper objectMapper;
    @Resource
    private AiMcpPayloadMapper aiMcpPayloadMapper;

    /**
     * 处理 MCP 协议入口请求，兼容批量与单条 JSON-RPC 负载。
     */
    public Object handle(Long tenantId, JsonNode body) {
        if (body != null && body.isArray()) {
            List<Object> responses = new ArrayList<>();
            for (JsonNode item : body) {
                Object response = handleSingle(tenantId, item);
                if (response != null) {
                    responses.add(response);
                }
            }
            return responses;
        }
        return handleSingle(tenantId, body);
    }

    /**
     * 处理单条 JSON-RPC MCP 请求。
     */
    public Object handleSingle(Long tenantId, JsonNode request) {
        if (request == null || request.isMissingNode() || request.isNull()) {
            return error(null, -32600, "invalid request");
        }
        JsonNode idNode = request.get("id");
        String method = request.path("method").asText("");
        try {
            if ("initialize".equals(method)) {
                return success(idNode, buildInitializeResult());
            }
            if ("notifications/initialized".equals(method)) {
                return idNode == null || idNode.isNull() ? null : success(idNode, new LinkedHashMap<String, Object>());
            }
            if ("ping".equals(method)) {
                return success(idNode, new LinkedHashMap<String, Object>());
            }
            if ("tools/list".equals(method)) {
                return success(idNode, buildToolsListResult(tenantId));
            }
            if ("tools/call".equals(method)) {
                return success(idNode, buildToolCallResult(tenantId, request.path("params")));
            }
            return error(idNode, -32601, "method not found");
        } catch (Exception e) {
            return error(idNode, -32000, e.getMessage());
        }
    }

    /**
     * 构造 initialize 方法的标准返回体。
     */
    private Map<String, Object> buildInitializeResult() {
        Map<String, Object> result = new LinkedHashMap<>();
        result.put("protocolVersion", AiMcpConstants.PROTOCOL_VERSION);
        Map<String, Object> capabilities = new LinkedHashMap<>();
        Map<String, Object> tools = new LinkedHashMap<>();
        tools.put("listChanged", false);
        capabilities.put("tools", tools);
        result.put("capabilities", capabilities);
        Map<String, Object> serverInfo = new LinkedHashMap<>();
        serverInfo.put("name", AiMcpConstants.SERVER_NAME);
        serverInfo.put("version", AiMcpConstants.SERVER_VERSION);
        result.put("serverInfo", serverInfo);
        return result;
    }

    /**
     * 构造 tools/list 方法返回体。
     */
    private Map<String, Object> buildToolsListResult(Long tenantId) {
        List<Map<String, Object>> tools = new ArrayList<>();
        for (AiMcpToolDescriptor descriptor : aiMcpRegistryService.listTools(tenantId, null)) {
            if (descriptor == null || !Integer.valueOf(1).equals(descriptor.getEnabledFlag())) {
                continue;
            }
            tools.add(aiMcpPayloadMapper.toProtocolTool(descriptor));
        }
        Map<String, Object> result = new LinkedHashMap<>();
        result.put("tools", tools);
        return result;
    }

    /**
     * 执行 tools/call，并把内部工具结果转换成 MCP 协议输出。
     */
    private Map<String, Object> buildToolCallResult(Long tenantId, JsonNode paramsNode) {
        String name = paramsNode.path("name").asText("");
        Map<String, Object> arguments = paramsNode.has("arguments") && !paramsNode.get("arguments").isNull()
                ? objectMapper.convertValue(paramsNode.get("arguments"), Map.class)
                : new LinkedHashMap<String, Object>();
        AiPromptContext context = new AiPromptContext()
                .setTenantId(tenantId)
                .setSceneCode(arguments.get("sceneCode") == null ? "system_diagnose" : String.valueOf(arguments.get("sceneCode")))
                .setQuestion(arguments.get("question") == null ? "请执行一次MCP工具调用" : String.valueOf(arguments.get("question")));
        AiDiagnosticToolResult result = aiMcpRegistryService.executeTool(tenantId, name, context, arguments);
        Map<String, Object> payload = new LinkedHashMap<>();
        List<Map<String, Object>> content = new ArrayList<>();
        Map<String, Object> text = new LinkedHashMap<>();
        text.put("type", "text");
        text.put("text", result == null || result.getSummaryText() == null ? "" : result.getSummaryText());
        content.add(text);
        payload.put("content", content);
        payload.put("isError", result != null && "WARN".equalsIgnoreCase(result.getSeverity()));
        if (result != null && result.getRawMeta() != null && !result.getRawMeta().isEmpty()) {
            payload.put("structuredContent", result.getRawMeta());
        }
        return payload;
    }

    /**
     * 生成 JSON-RPC 成功响应。
     */
    public Map<String, Object> success(JsonNode idNode, Object result) {
        Map<String, Object> payload = new LinkedHashMap<>();
        payload.put("jsonrpc", "2.0");
        payload.put("id", idNode == null || idNode.isNull() ? null : objectMapper.convertValue(idNode, Object.class));
        payload.put("result", result == null ? new LinkedHashMap<String, Object>() : result);
        return payload;
    }

    /**
     * 生成 JSON-RPC 错误响应。
     */
    public Map<String, Object> error(JsonNode idNode, int code, String message) {
        Map<String, Object> payload = new LinkedHashMap<>();
        payload.put("jsonrpc", "2.0");
        payload.put("id", idNode == null || idNode.isNull() ? null : objectMapper.convertValue(idNode, Object.class));
        Map<String, Object> error = new LinkedHashMap<>();
        error.put("code", code);
        error.put("message", message == null ? "unknown error" : message);
        payload.put("error", error);
        return payload;
    }
}