/*
 * Copyright 2008-2026 Async-IO.org
 *
 * Licensed under the Apache License, Version 2.0 (the "License"); you may not
 * use this file except in compliance with the License. You may obtain a copy of
 * the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
 * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
 * License for the specific language governing permissions and limitations under
 * the License.
 */
package org.atmosphere.mcp.runtime;

import org.atmosphere.cpr.AtmosphereHandler;
import org.atmosphere.cpr.AtmosphereResource;
import org.atmosphere.cpr.AtmosphereResourceEvent;
import org.atmosphere.cpr.AtmosphereResponse;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import tools.jackson.databind.ObjectMapper;

import java.io.IOException;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;

/**
 * {@link AtmosphereHandler} that bridges Atmosphere's transport layer with the
 * MCP JSON-RPC protocol. Supports three transports:
 * <ul>
 *   <li><b>Streamable HTTP</b> (MCP spec 2025-11-25) — POST for requests, GET for SSE notifications, DELETE to end session</li>
 *   <li><b>WebSocket</b> — full-duplex JSON-RPC via WebSocket frames</li>
 *   <li><b>SSE fallback</b> — Atmosphere's automatic WebSocket→SSE downgrade</li>
 * </ul>
 * Sessions survive transient disconnections: pending notifications are buffered
 * and replayed when a client reconnects with the same {@code Mcp-Session-Id}.
 * Idle sessions are evicted after a configurable TTL (default 30 min).
 * <p>
 * Registered at the MCP path resolved from {@code @Agent}.
 */
public final class McpHandler implements AtmosphereHandler {

    private static final Logger logger = LoggerFactory.getLogger(McpHandler.class);
    private static final String APPLICATION_JSON = "application/json";
    private static final String TEXT_EVENT_STREAM = "text/event-stream";
    private static final int DEFAULT_MAX_SESSIONS = 10_000;
    private static final int MAX_BODY_SIZE = 10 * 1024 * 1024;
    private static final ObjectMapper MAPPER = new ObjectMapper();

    /**
     * Header name for MCP protocol version negotiation, required on all
     * post-handshake requests since spec 2025-06-18.
     */
    public static final String PROTOCOL_VERSION_HEADER = "MCP-Protocol-Version";

    private final McpProtocolHandler protocolHandler;
    private final Map<String, McpSession> sessions = new ConcurrentHashMap<>();
    private final long sessionTtlMs;
    private final int maxSessions;
    private final ScheduledExecutorService cleaner;

    public McpHandler(McpProtocolHandler protocolHandler) {
        this(protocolHandler, McpSession.DEFAULT_TTL_MS, DEFAULT_MAX_SESSIONS);
    }

    public McpHandler(McpProtocolHandler protocolHandler, long sessionTtlMs) {
        this(protocolHandler, sessionTtlMs, DEFAULT_MAX_SESSIONS);
    }

    public McpHandler(McpProtocolHandler protocolHandler, long sessionTtlMs, int maxSessions) {
        this.protocolHandler = protocolHandler;
        this.sessionTtlMs = sessionTtlMs;
        this.maxSessions = maxSessions;
        this.cleaner = Executors.newSingleThreadScheduledExecutor(r -> {
            var t = Thread.ofPlatform().daemon().name("mcp-session-cleaner").unstarted(r);
            return t;
        });
        this.cleaner.scheduleAtFixedRate(this::evictExpiredSessions,
                sessionTtlMs, Math.max(sessionTtlMs / 2, 60_000L), TimeUnit.MILLISECONDS);
    }

    @Override
    public void onRequest(AtmosphereResource resource) throws IOException {
        var request = resource.getRequest();
        var response = resource.getResponse();

        // MCP 2025-11-25 clarification: invalid Origin → HTTP 403, not 400.
        // We treat "no Origin" as benign (non-browser clients); a present
        // but non-matching Origin against a configured allowlist is the
        // case the spec is targeting. Allowlist lives on the request
        // attribute "org.atmosphere.mcp.allowedOrigins" (List<String>);
        // when absent, no enforcement (back-compat with existing deployments).
        if (!isOriginAllowed(request)) {
            response.setStatus(403);
            response.setContentType(APPLICATION_JSON);
            response.getWriter().write(
                    "{\"jsonrpc\":\"2.0\",\"error\":{\"code\":-32600,\"message\":\"Forbidden Origin\"}}");
            return;
        }

        var method = request.getMethod();

        switch (method.toUpperCase()) {
            case "POST" -> handlePost(resource);
            case "GET" -> handleGet(resource);
            case "DELETE" -> handleDelete(resource);
            default -> {
                response.setStatus(405);
                response.getWriter().write("{\"error\":\"Method not allowed\"}");
            }
        }
    }

    /**
     * Validate the {@code Origin} header against an optional allowlist.
     * Returns {@code true} when there is no allowlist configured, when no
     * {@code Origin} header is present (non-browser clients), or when the
     * presented Origin is in the allowlist.
     */
    @SuppressWarnings("unchecked")
    private boolean isOriginAllowed(org.atmosphere.cpr.AtmosphereRequest request) {
        var origin = request.getHeader("Origin");
        if (origin == null || origin.isEmpty()) {
            return true;
        }
        var allowed = (List<String>) request.getAttribute("org.atmosphere.mcp.allowedOrigins");
        if (allowed == null || allowed.isEmpty()) {
            return true;
        }
        return allowed.contains(origin);
    }

    /**
     * After the initialize handshake, MCP 2025-06-18 requires the client to
     * echo the negotiated protocol revision in {@code MCP-Protocol-Version}
     * on every subsequent HTTP request. A missing or mismatched header on a
     * known session yields HTTP 400 so misconfigured clients fail loud.
     *
     * @return {@code true} if the request may proceed, {@code false} if a
     *         response has already been written (caller should return).
     */
    private boolean enforceProtocolVersionHeader(AtmosphereResource resource, McpSession session) throws IOException {
        if (session == null || !session.isInitialized()) {
            // Pre-handshake (initialize itself) — no enforcement
            return true;
        }
        var negotiated = session.protocolVersion();
        if (negotiated == null) {
            return true;
        }
        var presented = resource.getRequest().getHeader(PROTOCOL_VERSION_HEADER);
        if (presented == null) {
            // Tolerant for back-compat: spec says clients SHOULD send it; we
            // don't fail-close here to avoid breaking existing 2025-03-26 clients
            // that don't know about the header. Once this is widely adopted we
            // can promote to fail-closed.
            return true;
        }
        if (!presented.equals(negotiated)) {
            var response = resource.getResponse();
            response.setStatus(400);
            response.setContentType(APPLICATION_JSON);
            // JSON-encode via Jackson so a malicious header value (e.g. one
            // containing `"` or backslashes) cannot break out of the message
            // string and inject extra JSON fields. CodeQL flags raw string
            // concatenation of user-controlled headers (java/xss).
            response.getWriter().write(MAPPER.writeValueAsString(Map.of(
                    "jsonrpc", "2.0",
                    "error", Map.of(
                            "code", -32600,
                            "message", "MCP-Protocol-Version mismatch: session="
                                    + negotiated + ", header=" + presented))));
            return false;
        }
        return true;
    }

    /**
     * POST — Streamable HTTP: client sends JSON-RPC request, server responds with JSON or SSE.
     */
    private void handlePost(AtmosphereResource resource) throws IOException {
        var request = resource.getRequest();
        var response = resource.getResponse();

        // Reject oversized bodies early via Content-Length
        var contentLengthStr = request.getHeader("Content-Length");
        if (contentLengthStr != null) {
            try {
                long contentLength = Long.parseLong(contentLengthStr);
                if (contentLength > MAX_BODY_SIZE) {
                    response.setStatus(413);
                    response.setContentType(APPLICATION_JSON);
                    response.getWriter().write(
                            "{\"jsonrpc\":\"2.0\",\"error\":{\"code\":-32600,\"message\":\"Request body too large\"}}");
                    return;
                }
            } catch (NumberFormatException ignored) {
                // Malformed Content-Length; bounded read below will protect
            }
        }

        // Read request body with size limit
        var reader = request.getReader();
        var sb = new StringBuilder();
        var buf = new char[4096];
        int totalRead = 0;
        int n;
        while ((n = reader.read(buf)) != -1) {
            totalRead += n;
            if (totalRead > MAX_BODY_SIZE) {
                response.setStatus(413);
                response.setContentType(APPLICATION_JSON);
                response.getWriter().write(
                        "{\"jsonrpc\":\"2.0\",\"error\":{\"code\":-32600,\"message\":\"Request body too large\"}}");
                return;
            }
            sb.append(buf, 0, n);
        }

        if (sb.isEmpty()) {
            response.setStatus(400);
            response.setContentType(APPLICATION_JSON);
            response.getWriter().write("{\"jsonrpc\":\"2.0\",\"error\":{\"code\":-32700,\"message\":\"Empty body\"}}");
            return;
        }

        // Restore session from Mcp-Session-Id header if present
        var sessionId = request.getHeader(McpSession.SESSION_ID_HEADER);
        McpSession existingSession = null;
        if (sessionId != null) {
            existingSession = sessions.get(sessionId);
            if (existingSession != null) {
                request.setAttribute(McpSession.ATTRIBUTE_KEY, existingSession);
            }
        }

        // MCP 2025-06-18: protocol version header enforcement on post-handshake POSTs
        if (!enforceProtocolVersionHeader(resource, existingSession)) {
            return;
        }

        var jsonResponse = protocolHandler.handleMessage(resource, sb.toString());

        // After handleMessage, check if a new session was created (initialize)
        var session = (McpSession) request.getAttribute(McpSession.ATTRIBUTE_KEY);
        if (session != null) {
            try {
                sessions.compute(session.sessionId(), (id, existing) -> {
                    if (existing != null) return existing;
                    if (sessions.size() >= maxSessions) {
                        throw new IllegalStateException("Too many sessions");
                    }
                    return session;
                });
            } catch (IllegalStateException e) {
                response.setStatus(503);
                response.setContentType(APPLICATION_JSON);
                response.getWriter().write(
                        "{\"jsonrpc\":\"2.0\",\"error\":{\"code\":-32000,\"message\":\"Too many sessions\"}}");
                return;
            }
            session.touch();
            response.setHeader(McpSession.SESSION_ID_HEADER, session.sessionId());
        }

        if (jsonResponse == null) {
            // Notification (no id) — respond with 202 Accepted
            response.setStatus(202);
            return;
        }

        var accept = request.getHeader("Accept");
        if (accept != null && accept.contains(TEXT_EVENT_STREAM)) {
            // SSE format
            response.setStatus(200);
            response.setContentType(TEXT_EVENT_STREAM);
            response.setCharacterEncoding("UTF-8");
            response.getWriter().write("event: message\ndata: " + jsonResponse + "\n\n");
            response.getWriter().flush();
        } else {
            // Plain JSON
            response.setStatus(200);
            response.setContentType(APPLICATION_JSON);
            response.getWriter().write(jsonResponse);
            response.getWriter().flush();
        }
    }

    /**
     * GET — Opens SSE stream for server-initiated notifications.
     */
    private void handleGet(AtmosphereResource resource) throws IOException {
        var request = resource.getRequest();
        var response = resource.getResponse();

        // Restore session from header
        var sessionId = request.getHeader(McpSession.SESSION_ID_HEADER);
        McpSession session = null;
        if (sessionId != null) {
            session = sessions.get(sessionId);
            if (session != null) {
                request.setAttribute(McpSession.ATTRIBUTE_KEY, session);
                session.touch();
                response.setHeader(McpSession.SESSION_ID_HEADER, session.sessionId());
            }
        }
        if (!enforceProtocolVersionHeader(resource, session)) {
            return;
        }

        // SSE notification stream or WebSocket upgrade
        response.setContentType(TEXT_EVENT_STREAM);
        response.setCharacterEncoding("UTF-8");
        resource.suspend();

        // Replay any notifications buffered while the client was disconnected
        if (session != null) {
            var pending = session.drainPendingNotifications();
            for (var notification : pending) {
                response.getWriter().write("event: message\ndata: " + notification + "\n\n");
            }
            if (!pending.isEmpty()) {
                response.getWriter().flush();
                logger.debug("Replayed {} pending notifications for session {}", pending.size(), sessionId);
            }
        }
    }

    /**
     * DELETE — Terminates MCP session.
     */
    private void handleDelete(AtmosphereResource resource) throws IOException {
        var request = resource.getRequest();
        var response = resource.getResponse();
        var sessionId = request.getHeader(McpSession.SESSION_ID_HEADER);

        McpSession session = null;
        if (sessionId != null) {
            session = sessions.get(sessionId);
        }
        if (!enforceProtocolVersionHeader(resource, session)) {
            return;
        }

        if (sessionId != null) {
            var removed = sessions.remove(sessionId);
            if (removed != null) {
                logger.info("MCP session terminated: {}", sessionId);
            }
        }

        response.setStatus(204);
    }

    @Override
    public void onStateChange(AtmosphereResourceEvent event) throws IOException {
        var resource = event.getResource();
        if (event.isCancelled() || event.isClosedByClient() || event.isClosedByApplication()) {
            logger.debug("MCP connection closed: {}", resource.uuid());
            return;
        }

        var message = event.getMessage();
        if (message instanceof String text) {
            // Incoming WebSocket message
            var response = protocolHandler.handleMessage(resource, text);
            if (response != null) {
                write(resource.getResponse(), response);
            }
        } else if (message instanceof List<?> list) {
            // Broadcast — write each message
            for (var item : list) {
                if (item instanceof String text) {
                    write(resource.getResponse(), text);
                }
            }
        }
    }

    /**
     * Called directly for WebSocket messages (bypasses onStateChange for request-response).
     */
    public void processMessage(AtmosphereResource resource, String message) throws IOException {
        var response = protocolHandler.handleMessage(resource, message);
        if (response != null) {
            write(resource.getResponse(), response);
        }
    }

    /**
     * Returns the session store (visible for testing).
     */
    public Map<String, McpSession> sessions() {
        return sessions;
    }

    @Override
    public void destroy() {
        cleaner.shutdownNow();
        sessions.clear();
        logger.debug("McpHandler destroyed");
    }

    private void evictExpiredSessions() {
        var it = sessions.entrySet().iterator();
        while (it.hasNext()) {
            var entry = it.next();
            if (entry.getValue().isExpired(sessionTtlMs)) {
                it.remove();
                logger.info("Evicted expired MCP session: {}", entry.getKey());
            }
        }
    }

    private void write(AtmosphereResponse response, String data) throws IOException {
        response.getWriter().write(data);
        response.getWriter().flush();
    }
}
