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.system.entity.AiMcpMount;
import org.springframework.stereotype.Component;

import javax.annotation.Resource;
import java.io.BufferedReader;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.io.OutputStream;
import java.net.HttpURLConnection;
import java.net.URL;
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;

@Component
public class AiMcpHttpClient {

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

    /**
     * 通过 Streamable HTTP 协议加载远程 MCP 工具目录。
     */
    public List<AiMcpToolDescriptor> listTools(AiMcpMount mount) {
        initialize(mount);
        JsonNode result = sendRequest(mount, "tools/list", new LinkedHashMap<String, Object>(), true);
        List<AiMcpToolDescriptor> output = new ArrayList<>();
        JsonNode toolsNode = result.path("tools");
        if (!toolsNode.isArray()) {
            return output;
        }
        for (JsonNode item : toolsNode) {
            AiMcpToolDescriptor descriptor = aiMcpPayloadMapper.toExternalToolDescriptor(mount, item);
            if (descriptor != null) {
                output.add(descriptor);
            }
        }
        return output;
    }

    /**
     * 通过 Streamable HTTP 协议执行一次远程工具调用。
     */
    public AiDiagnosticToolResult callTool(AiMcpMount mount, String toolName, Map<String, Object> arguments) {
        initialize(mount);
        Map<String, Object> params = new LinkedHashMap<>();
        params.put("name", toolName);
        params.put("arguments", arguments == null ? new LinkedHashMap<String, Object>() : arguments);
        JsonNode result = sendRequest(mount, "tools/call", params, true);
        return aiMcpPayloadMapper.toExternalToolResult(mount, toolName, result);
    }

    /**
     * 执行 MCP initialize + notifications/initialized 握手。
     */
    private void initialize(AiMcpMount mount) {
        Map<String, Object> params = new LinkedHashMap<>();
        params.put("protocolVersion", AiMcpConstants.PROTOCOL_VERSION);
        params.put("capabilities", new LinkedHashMap<String, Object>());
        Map<String, Object> clientInfo = new LinkedHashMap<>();
        clientInfo.put("name", "rsf-server");
        clientInfo.put("version", AiMcpConstants.SERVER_VERSION);
        params.put("clientInfo", clientInfo);
        sendRequest(mount, "initialize", params, true);
        sendRequest(mount, "notifications/initialized", new LinkedHashMap<String, Object>(), false);
    }

    /**
     * 发送一条 JSON-RPC 请求到远程 MCP HTTP 端点。
     */
    private JsonNode sendRequest(AiMcpMount mount, String method, Object params, boolean expectResponse) {
        HttpURLConnection connection = null;
        try {
            connection = (HttpURLConnection) new URL(mount.getUrl()).openConnection();
            connection.setRequestMethod("POST");
            connection.setDoOutput(true);
            connection.setConnectTimeout(mount.getTimeoutMs() == null ? 10000 : mount.getTimeoutMs());
            connection.setReadTimeout(mount.getTimeoutMs() == null ? 10000 : mount.getTimeoutMs());
            connection.setRequestProperty("Content-Type", "application/json");
            connection.setRequestProperty("Accept", "application/json");
            applyAuthHeaders(connection, mount);

            Map<String, Object> body = new LinkedHashMap<>();
            body.put("jsonrpc", "2.0");
            if (expectResponse) {
                body.put("id", String.valueOf(System.currentTimeMillis()));
            }
            body.put("method", method);
            body.put("params", params == null ? new LinkedHashMap<String, Object>() : params);

            try (OutputStream outputStream = connection.getOutputStream()) {
                outputStream.write(objectMapper.writeValueAsBytes(body));
                outputStream.flush();
            }

            int statusCode = connection.getResponseCode();
            InputStream inputStream = statusCode >= 400 ? connection.getErrorStream() : connection.getInputStream();
            if (!expectResponse) {
                return null;
            }
            if (inputStream == null) {
                throw new IllegalStateException("MCP服务返回空响应");
            }
            String payload = readPayload(inputStream);
            JsonNode root = objectMapper.readTree(payload);
            if (root.has("error") && !root.get("error").isNull()) {
                throw new IllegalStateException(root.path("error").path("message").asText("MCP调用失败"));
            }
            return root.path("result");
        } catch (Exception e) {
            throw new IllegalStateException("MCP请求失败: " + e.getMessage(), e);
        } finally {
            if (connection != null) {
                connection.disconnect();
            }
        }
    }

    /**
     * 按挂载配置写入鉴权请求头。
     */
    private void applyAuthHeaders(HttpURLConnection connection, AiMcpMount mount) {
        if (mount == null || mount.getAuthType() == null || mount.getAuthValue() == null || mount.getAuthValue().trim().isEmpty()) {
            return;
        }
        String authType = mount.getAuthType().trim().toUpperCase();
        if (AiMcpConstants.AUTH_TYPE_BEARER.equals(authType)) {
            connection.setRequestProperty("Authorization", "Bearer " + mount.getAuthValue().trim());
        } else if (AiMcpConstants.AUTH_TYPE_API_KEY.equals(authType)) {
            connection.setRequestProperty("X-API-Key", mount.getAuthValue().trim());
        }
    }

    /**
     * 读取 HTTP 响应体全文。
     */
    private String readPayload(InputStream inputStream) throws Exception {
        StringBuilder output = new StringBuilder();
        try (BufferedReader reader = new BufferedReader(new InputStreamReader(inputStream, StandardCharsets.UTF_8))) {
            String line;
            while ((line = reader.readLine()) != null) {
                output.append(line);
            }
        }
        return output.toString();
    }
}