package com.gzzm.lobster.tool.mcp;

import com.gzzm.lobster.common.JsonUtil;
import com.gzzm.lobster.common.LobsterException;
import com.gzzm.lobster.common.ToolRiskLevel;

import java.io.ByteArrayOutputStream;
import java.io.InputStream;
import java.io.OutputStream;
import java.net.HttpURLConnection;
import java.net.URL;
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.Collections;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;

/**
 * 调用 API Center 暴露的 MCP JSON-RPC HTTP 入口。
 */
public class HttpApiCenterMcpClient implements McpClient {

    private static final String MCP_PROTOCOL_VERSION = "2025-11-25";

    @Override
    public List<McpToolSpec> listTools(McpServerConfig server, McpCallOptions options) throws Exception {
        List<McpToolSpec> out = new ArrayList<>();
        String cursor = null;
        int guard = 0;
        do {
            Map<String, Object> params = new LinkedHashMap<>();
            params.put("limit", Integer.valueOf(100));
            if (cursor != null) params.put("cursor", cursor);
            Map<String, Object> result = jsonRpc(server, "tools/list", params, options);
            out.addAll(McpClientSupport.parseToolList(result, defaultRisk(server)));
            Object next = result.get("nextCursor");
            cursor = next == null ? null : String.valueOf(next);
            guard++;
        } while (cursor != null && !cursor.isEmpty() && guard < 20);
        return out;
    }

    @Override
    public McpCallResult callTool(McpServerConfig server,
                                  String remoteToolName,
                                  Map<String, Object> args,
                                  McpCallOptions options) throws Exception {
        Map<String, Object> params = new LinkedHashMap<>();
        params.put("name", remoteToolName);
        params.put("arguments", args == null ? Collections.<String, Object>emptyMap() : args);
        Map<String, Object> result = jsonRpc(server, "tools/call", params, options);
        return McpClientSupport.parseCallResult(result);
    }

    @SuppressWarnings("unchecked")
    private Map<String, Object> jsonRpc(McpServerConfig server,
                                        String method,
                                        Map<String, Object> params,
                                        McpCallOptions options) throws Exception {
        if (server == null || server.getEndpoint() == null || server.getEndpoint().trim().isEmpty()) {
            throw new LobsterException("mcp.endpoint", "MCP endpoint is required");
        }
        Map<String, Object> body = new LinkedHashMap<>();
        body.put("jsonrpc", "2.0");
        body.put("id", options == null || options.getRequestId() == null ? "lobster-mcp" : options.getRequestId());
        body.put("method", method);
        body.put("params", params == null ? Collections.<String, Object>emptyMap() : params);
        RpcEnvelope envelope = postJson(server, body, options);
        Map<String, Object> response = envelope.body;
        if (envelope.status >= 400) {
            throw new LobsterException("mcp.http_error",
                    "MCP HTTP error " + envelope.status + " from " + server.getEndpoint());
        }
        Object error = response.get("error");
        if (error instanceof Map) {
            Object message = ((Map<?, ?>) error).get("message");
            throw new LobsterException("mcp.remote_error", message == null ? "MCP remote error" : String.valueOf(message));
        }
        Object result = response.get("result");
        if (response.containsKey("result") && result instanceof Map) {
            return (Map<String, Object>) result;
        }
        throw new LobsterException("mcp.invalid_response",
                "MCP response missing object result for method " + method);
    }

    @SuppressWarnings("unchecked")
    private RpcEnvelope postJson(McpServerConfig server,
                                 Map<String, Object> body,
                                 McpCallOptions options) throws Exception {
        HttpURLConnection conn = null;
        try {
            conn = (HttpURLConnection) new URL(server.getEndpoint()).openConnection();
            conn.setRequestMethod("POST");
            conn.setDoOutput(true);
            conn.setConnectTimeout(timeout(server.getConnectTimeoutMs(), 5000));
            conn.setReadTimeout(timeout(server.getReadTimeoutMs(), optionTimeout(options, 30000)));
            conn.setRequestProperty("Content-Type", "application/json; charset=UTF-8");
            conn.setRequestProperty("Accept", "application/json, text/event-stream");
            conn.setRequestProperty("MCP-Protocol-Version", MCP_PROTOCOL_VERSION);
            String requestId = options == null ? null : options.getRequestId();
            if (requestId != null) conn.setRequestProperty("X-Request-Id", requestId);
            applyConfiguredHeaders(conn, server);
            if (server.getAuthToken() != null && !server.getAuthToken().trim().isEmpty()) {
                conn.setRequestProperty("X-ZGW-Auth", "Bearer " + server.getAuthToken().trim());
            }
            if (options != null) {
                if (options.getThreadId() != null) conn.setRequestProperty("X-MCP-Lobster-Thread-Id", options.getThreadId());
                if (options.getRunId() != null) conn.setRequestProperty("X-MCP-Lobster-Run-Id", options.getRunId());
                if (options.getToolCallId() != null) conn.setRequestProperty("X-MCP-Lobster-Tool-Call-Id", options.getToolCallId());
            }
            byte[] bytes = JsonUtil.toJson(body).getBytes(StandardCharsets.UTF_8);
            try (OutputStream os = conn.getOutputStream()) {
                os.write(bytes);
            }
            int status = conn.getResponseCode();
            String text = readBody(status >= 400 ? conn.getErrorStream() : conn.getInputStream());
            if (text == null || text.trim().isEmpty()) {
                if (status >= 400) {
                    throw new LobsterException("mcp.http_error",
                            "MCP HTTP error " + status + " from " + server.getEndpoint());
                }
                throw new LobsterException("mcp.invalid_response", "MCP response body is empty");
            }
            Object parsed;
            try {
                parsed = JsonUtil.fromJson(text, Object.class);
            } catch (Throwable parseError) {
                throw new LobsterException(status >= 400 ? "mcp.http_error" : "mcp.invalid_response",
                        "MCP response is not JSON: HTTP " + status + " from " + server.getEndpoint(), parseError);
            }
            if (parsed instanceof Map) {
                return new RpcEnvelope(status, (Map<String, Object>) parsed);
            }
            throw new LobsterException(status >= 400 ? "mcp.http_error" : "mcp.invalid_response",
                    "MCP response JSON root is not an object: HTTP " + status + " from " + server.getEndpoint());
        } finally {
            if (conn != null) conn.disconnect();
        }
    }

    @SuppressWarnings("unchecked")
    private void applyConfiguredHeaders(HttpURLConnection conn, McpServerConfig server) {
        String raw = server.getHeadersJson();
        if (raw == null || raw.trim().isEmpty()) return;
        try {
            Object parsed = JsonUtil.fromJson(raw, Object.class);
            if (!(parsed instanceof Map)) return;
            for (Map.Entry<?, ?> e : ((Map<?, ?>) parsed).entrySet()) {
                if (e.getKey() != null && e.getValue() != null) {
                    conn.setRequestProperty(String.valueOf(e.getKey()), String.valueOf(e.getValue()));
                }
            }
        } catch (Throwable ignore) { /* invalid custom headers are ignored */ }
    }

    private String readBody(InputStream in) throws Exception {
        if (in == null) return "";
        try (InputStream input = in; ByteArrayOutputStream out = new ByteArrayOutputStream()) {
            byte[] buf = new byte[8192];
            int n;
            while ((n = input.read(buf)) >= 0) {
                out.write(buf, 0, n);
            }
            return new String(out.toByteArray(), StandardCharsets.UTF_8);
        }
    }

    private int timeout(Integer value, int fallback) {
        return value == null || value <= 0 ? fallback : value;
    }

    private int optionTimeout(McpCallOptions options, int fallback) {
        if (options == null || options.getTimeoutMs() <= 0) return fallback;
        return options.getTimeoutMs() > Integer.MAX_VALUE
                ? Integer.MAX_VALUE : (int) options.getTimeoutMs();
    }

    private ToolRiskLevel defaultRisk(McpServerConfig server) {
        return server == null || server.getDefaultRisk() == null ? ToolRiskLevel.READ_ONLY : server.getDefaultRisk();
    }

    private static final class RpcEnvelope {
        final int status;
        final Map<String, Object> body;

        RpcEnvelope(int status, Map<String, Object> body) {
            this.status = status;
            this.body = body == null ? Collections.<String, Object>emptyMap() : body;
        }
    }
}