package com.smqjh.agent.mcp;

import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.node.ArrayNode;
import com.fasterxml.jackson.databind.node.ObjectNode;
import com.sun.net.httpserver.Headers;
import com.sun.net.httpserver.HttpExchange;
import com.sun.net.httpserver.HttpServer;

import java.io.IOException;
import java.net.InetSocketAddress;
import java.nio.charset.StandardCharsets;
import java.time.Instant;
import java.util.Objects;
import java.util.concurrent.Executors;

final class McpHttpServer {
    private final AppConfig config;
    private final ToolRegistry tools;
    private HttpServer server;

    McpHttpServer(AppConfig config, ToolRegistry tools) {
        this.config = config;
        this.tools = tools;
    }

    void start() throws IOException {
        if (!config.mcp.enabled) {
            System.out.println("MCP disabled");
            return;
        }
        String host = config.mcp.allowExternal ? config.mcp.host : "127.0.0.1";
        server = HttpServer.create(new InetSocketAddress(host, config.mcp.port), 0);
        server.createContext(normalizedPath(), this::handle);
        server.setExecutor(Executors.newFixedThreadPool(8));
        server.start();
        System.out.printf("SMQJH MCP server started: http://%s:%d%s%n", host, config.mcp.port, normalizedPath());
    }

    void stop() {
        if (server != null) {
            server.stop(0);
            System.out.println("SMQJH MCP server stopped");
        }
    }

    private void handle(HttpExchange exchange) throws IOException {
        applyCors(exchange.getResponseHeaders());
        try {
            if ("OPTIONS".equalsIgnoreCase(exchange.getRequestMethod())) {
                send(exchange, 204, "");
                return;
            }
            if (!authorized(exchange)) {
                sendJson(exchange, 401, Jsons.object().put("ok", false).put("message", "MCP token required"));
                return;
            }
            if ("GET".equalsIgnoreCase(exchange.getRequestMethod())) {
                sendJson(exchange, 200, statusPayload());
                return;
            }
            if (!"POST".equalsIgnoreCase(exchange.getRequestMethod())) {
                sendJson(exchange, 405, Jsons.object().put("ok", false).put("message", "Only GET and POST are supported"));
                return;
            }

            byte[] bodyBytes = exchange.getRequestBody().readAllBytes();
            if (bodyBytes.length > 1024 * 1024) {
                sendJson(exchange, 413, rpcError(null, -32600, "Request body is too large"));
                return;
            }
            JsonNode request = bodyBytes.length == 0 ? Jsons.object() : Jsons.MAPPER.readTree(bodyBytes);
            if (request.isArray()) {
                ArrayNode replies = Jsons.MAPPER.createArrayNode();
                for (JsonNode item : request) {
                    JsonNode reply = handleRpc(item);
                    if (reply != null) {
                        replies.add(reply);
                    }
                }
                sendJson(exchange, 200, replies);
            } else {
                JsonNode reply = handleRpc(request);
                if (reply == null) {
                    send(exchange, 204, "");
                } else {
                    sendJson(exchange, 200, reply);
                }
            }
        } catch (Exception error) {
            sendJson(exchange, 500, rpcError(null, -32603, error.getMessage()));
        }
    }

    private JsonNode handleRpc(JsonNode request) {
        JsonNode idNode = request.path("id");
        String method = request.path("method").asText("");
        if ((idNode.isMissingNode() || idNode.isNull()) && method.startsWith("notifications/")) {
            return null;
        }
        try {
            return switch (method) {
                case "initialize" -> rpcResult(idNode, initializePayload());
                case "tools/list" -> rpcResult(idNode, Jsons.object().set("tools", Jsons.MAPPER.valueToTree(tools.listTools())));
                case "tools/call" -> rpcResult(idNode, tools.call(request.path("params")));
                default -> rpcError(idNode, -32601, "Unknown MCP method: " + method);
            };
        } catch (Exception error) {
            return rpcError(idNode, -32603, error.getMessage());
        }
    }

    private ObjectNode initializePayload() {
        ObjectNode capabilities = Jsons.object();
        capabilities.set("tools", Jsons.object());
        ObjectNode serverInfo = Jsons.object();
        serverInfo.put("name", "smqjh-java-mcp");
        serverInfo.put("version", "0.1.0");
        ObjectNode payload = Jsons.object();
        payload.put("protocolVersion", "2024-11-05");
        payload.set("capabilities", capabilities);
        payload.set("serverInfo", serverInfo);
        return payload;
    }

    private ObjectNode statusPayload() {
        ObjectNode payload = Jsons.object();
        payload.put("ok", true);
        payload.put("name", "smqjh-java-mcp");
        payload.put("environmentName", config.environmentName);
        payload.put("baseUrl", config.baseUrl);
        payload.put("startedAtCheck", Instant.now().toString());
        payload.put("path", normalizedPath());
        payload.put("toolCount", tools.listTools().size());
        payload.put("databaseEnabled", config.database.enabled);
        payload.put("desktopConfigPath", Objects.toString(config.desktopConfigPath, ""));
        return payload;
    }

    private ObjectNode rpcResult(JsonNode idNode, JsonNode result) {
        ObjectNode payload = Jsons.object();
        payload.put("jsonrpc", "2.0");
        putId(payload, idNode);
        payload.set("result", result);
        return payload;
    }

    private ObjectNode rpcError(JsonNode idNode, int code, String message) {
        ObjectNode error = Jsons.object();
        error.put("code", code);
        error.put("message", message == null ? "MCP error" : message);
        ObjectNode payload = Jsons.object();
        payload.put("jsonrpc", "2.0");
        putId(payload, idNode);
        payload.set("error", error);
        return payload;
    }

    private void putId(ObjectNode payload, JsonNode idNode) {
        if (idNode == null || idNode.isMissingNode() || idNode.isNull()) {
            payload.putNull("id");
        } else if (idNode.isNumber()) {
            payload.put("id", idNode.asLong());
        } else {
            payload.put("id", idNode.asText());
        }
    }

    private boolean authorized(HttpExchange exchange) {
        String token = config.mcp.authToken == null ? "" : config.mcp.authToken.trim();
        if (token.isEmpty()) {
            return true;
        }
        Headers headers = exchange.getRequestHeaders();
        String authorization = headers.getFirst("Authorization");
        String mcpToken = headers.getFirst("X-MCP-Token");
        return ("Bearer " + token).equals(authorization) || token.equals(mcpToken);
    }

    private void applyCors(Headers headers) {
        headers.set("Access-Control-Allow-Origin", "http://127.0.0.1");
        headers.set("Access-Control-Allow-Methods", "GET,POST,OPTIONS");
        headers.set("Access-Control-Allow-Headers", "Content-Type,Authorization,X-MCP-Token");
    }

    private String normalizedPath() {
        String path = config.mcp.path == null || config.mcp.path.isBlank() ? "/mcp" : config.mcp.path.trim();
        return path.startsWith("/") ? path : "/" + path;
    }

    private void sendJson(HttpExchange exchange, int status, JsonNode payload) throws IOException {
        exchange.getResponseHeaders().set("Content-Type", "application/json; charset=utf-8");
        send(exchange, status, Jsons.stringify(payload));
    }

    private void send(HttpExchange exchange, int status, String text) throws IOException {
        byte[] bytes = text.getBytes(StandardCharsets.UTF_8);
        exchange.sendResponseHeaders(status, bytes.length);
        exchange.getResponseBody().write(bytes);
        exchange.close();
    }
}
