/*
 * This software is in the public domain under CC0 1.0 Universal plus a 
 * Grant of Patent License.
 * 
 * To the extent possible under law, author(s) have dedicated all
 * copyright and related and neighboring rights to this software to the
 * public domain worldwide. This software is distributed without any
 * warranty.
 * 
 * You should have received a copy of the CC0 Public Domain Dedication
 * along with this software (see the LICENSE.md file). If not, see
 * <http://creativecommons.org/publicdomain/zero/1.0/>.
 */
package org.moqui.mcp

import groovy.json.JsonSlurper
import groovy.json.JsonOutput
import org.moqui.impl.context.ExecutionContextFactoryImpl
import org.moqui.context.ArtifactAuthorizationException
import org.moqui.context.ArtifactTarpitException
import org.moqui.impl.context.ExecutionContextImpl
import org.slf4j.Logger
import org.slf4j.LoggerFactory

import javax.servlet.ServletConfig
import javax.servlet.ServletException
import javax.servlet.http.HttpServlet
import javax.servlet.http.HttpServletRequest
import javax.servlet.http.HttpServletResponse
import java.util.concurrent.ConcurrentHashMap
import java.util.concurrent.atomic.AtomicBoolean
import java.util.UUID

/**
 * Simple JSON-RPC Message classes for MCP compatibility
 */
class JsonRpcMessage {
    String jsonrpc = "2.0"
}

class JsonRpcResponse extends JsonRpcMessage {
    Object id
    Object result
    Map error
    
    JsonRpcResponse(Object result, Object id) {
        this.result = result
        this.id = id
    }
    
    JsonRpcResponse(Map error, Object id) {
        this.error = error
        this.id = id
    }
    
    String toJson() {
        return JsonOutput.toJson(this)
    }
}

class JsonRpcNotification extends JsonRpcMessage {
    String method
    Object params
    
    JsonRpcNotification(String method, Object params = null) {
        this.method = method
        this.params = params
    }
    
    String toJson() {
        return JsonOutput.toJson(this)
    }
}

/**
 * Enhanced MCP Servlet with proper SSE handling inspired by HttpServletSseServerTransportProvider
 * This implementation provides better SSE support and session management.
 */
class EnhancedMcpServlet extends HttpServlet {
    protected final static Logger logger = LoggerFactory.getLogger(EnhancedMcpServlet.class)
    
    private JsonSlurper jsonSlurper = new JsonSlurper()
    
    // Session management using Moqui's Visit system directly
    // No need for separate session manager - Visit entity handles persistence
    
    @Override
    void init(ServletConfig config) throws ServletException {
        super.init(config)
        String webappName = config.getInitParameter("moqui-name") ?: 
            config.getServletContext().getInitParameter("moqui-name")
        logger.info("EnhancedMcpServlet initialized for webapp ${webappName}")
    }
    
    @Override
    void service(HttpServletRequest request, HttpServletResponse response) 
            throws ServletException, IOException {
        
        ExecutionContextFactoryImpl ecfi = 
            (ExecutionContextFactoryImpl) getServletContext().getAttribute("executionContextFactory")
        String webappName = getInitParameter("moqui-name") ?: 
            getServletContext().getInitParameter("moqui-name")
        
        if (ecfi == null || webappName == null) {
            response.sendError(HttpServletResponse.SC_INTERNAL_SERVER_ERROR, 
                "System is initializing, try again soon.")
            return
        }
        
        // Handle CORS
        if (handleCors(request, response, webappName, ecfi)) return
        
        long startTime = System.currentTimeMillis()
        
        if (logger.traceEnabled) {
            logger.trace("Start Enhanced MCP request to [${request.getPathInfo()}] at time [${startTime}] in session [${request.session.id}] thread [${Thread.currentThread().id}:${Thread.currentThread().name}]")
        }
        
        ExecutionContextImpl activeEc = ecfi.activeContext.get()
        if (activeEc != null) {
            logger.warn("In EnhancedMcpServlet.service there is already an ExecutionContext for user ${activeEc.user.username}")
            activeEc.destroy()
        }
        
        ExecutionContextImpl ec = ecfi.getEci()
        
try {
            // Handle Basic Authentication directly without triggering screen system
            String authzHeader = request.getHeader("Authorization")
            boolean authenticated = false
            
            if (authzHeader != null && authzHeader.length() > 6 && authzHeader.startsWith("Basic ")) {
                String basicAuthEncoded = authzHeader.substring(6).trim()
                String basicAuthAsString = new String(basicAuthEncoded.decodeBase64())
                int indexOfColon = basicAuthAsString.indexOf(":")
                if (indexOfColon > 0) {
                    String username = basicAuthAsString.substring(0, indexOfColon)
                    String password = basicAuthAsString.substring(indexOfColon + 1)
                    try {
                        ec.user.loginUser(username, password)
                        authenticated = true
                        logger.info("Enhanced MCP Basic auth successful for user: ${ec.user?.username}")
                    } catch (Exception e) {
                        logger.warn("Enhanced MCP Basic auth failed for user ${username}: ${e.message}")
                    }
                } else {
                    logger.warn("Enhanced MCP got bad Basic auth credentials string")
                }
            }
            
            // Re-enabled proper authentication - UserServices compilation issues resolved
            if (!authenticated || !ec.user?.userId) {
                logger.warn("Enhanced MCP authentication failed - no valid user authenticated")
                response.setStatus(HttpServletResponse.SC_UNAUTHORIZED)
                response.setContentType("application/json")
                response.setHeader("WWW-Authenticate", "Basic realm=\"Moqui MCP\"")
                response.writer.write(groovy.json.JsonOutput.toJson([
                    jsonrpc: "2.0",
                    error: [code: -32003, message: "Authentication required. Use Basic auth with valid Moqui credentials."],
                    id: null
                ]))
                return
            }
            
            // Route based on request method and path
            String requestURI = request.getRequestURI()
            String method = request.getMethod()
            
            if ("GET".equals(method) && requestURI.endsWith("/sse")) {
                handleSseConnection(request, response, ec)
            } else if ("POST".equals(method) && requestURI.endsWith("/message")) {
                handleMessage(request, response, ec)
            } else if ("POST".equals(method) && (requestURI.equals("/mcp") || requestURI.endsWith("/mcp"))) {
                // Handle POST requests to /mcp for JSON-RPC
                handleJsonRpc(request, response, ec)
            } else if ("GET".equals(method) && (requestURI.equals("/mcp") || requestURI.endsWith("/mcp"))) {
                // Handle GET requests to /mcp - maybe for server info or SSE fallback
                handleSseConnection(request, response, ec)
            } else {
                // Fallback to JSON-RPC handling
                handleJsonRpc(request, response, ec)
            }
            
        } catch (ArtifactAuthorizationException e) {
            logger.warn("Enhanced MCP Access Forbidden (no authz): " + e.message)
            response.setStatus(HttpServletResponse.SC_FORBIDDEN)
            response.setContentType("application/json")
            response.writer.write(groovy.json.JsonOutput.toJson([
                jsonrpc: "2.0",
                error: [code: -32001, message: "Access Forbidden: " + e.message],
                id: null
            ]))
        } catch (ArtifactTarpitException e) {
            logger.warn("Enhanced MCP Too Many Requests (tarpit): " + e.message)
            response.setStatus(429)
            if (e.getRetryAfterSeconds()) {
                response.addIntHeader("Retry-After", e.getRetryAfterSeconds())
            }
            response.setContentType("application/json")
            response.writer.write(groovy.json.JsonOutput.toJson([
                jsonrpc: "2.0",
                error: [code: -32002, message: "Too Many Requests: " + e.message],
                id: null
            ]))
        } catch (Throwable t) {
            logger.error("Error in Enhanced MCP request", t)
            response.setStatus(HttpServletResponse.SC_INTERNAL_SERVER_ERROR)
            response.setContentType("application/json")
            response.writer.write(groovy.json.JsonOutput.toJson([
                jsonrpc: "2.0",
                error: [code: -32603, message: "Internal error: " + t.message],
                id: null
            ]))
        } finally {
            ec.destroy()
        }
    }
    
    private void handleSseConnection(HttpServletRequest request, HttpServletResponse response, ExecutionContextImpl ec) 
            throws IOException {
        
        logger.info("Handling Enhanced SSE connection from ${request.remoteAddr}")
        
        // Enable async support for SSE
        if (request.isAsyncSupported()) {
            request.startAsync()
        }
        
        // Set SSE headers
        response.setContentType("text/event-stream")
        response.setCharacterEncoding("UTF-8")
        response.setHeader("Cache-Control", "no-cache")
        response.setHeader("Connection", "keep-alive")
        response.setHeader("Access-Control-Allow-Origin", "*")
        response.setHeader("X-Accel-Buffering", "no") // Disable nginx buffering
        
        // Get or create Visit (Moqui automatically creates Visit)
        def visit = ec.user.getVisit()
        if (!visit) {
            response.sendError(HttpServletResponse.SC_INTERNAL_SERVER_ERROR, "Failed to create Visit")
            return
        }
        
        // Create Visit-based session transport
        VisitBasedMcpSession session = new VisitBasedMcpSession(visit, response.writer, ec)
        
        try {
            // Send initial connection event
            def connectData = [
                type: "connected",
                sessionId: visit.visitId,
                timestamp: System.currentTimeMillis(),
                serverInfo: [
                    name: "Moqui MCP SSE Server",
                    version: "2.0.0",
                    protocolVersion: "2025-06-18",
                    architecture: "Visit-based sessions"
                ]
            ]
            sendSseEvent(response.writer, "connect", groovy.json.JsonOutput.toJson(connectData), 0)
            
            // Send endpoint info for message posting
            sendSseEvent(response.writer, "endpoint", "/mcp/message?sessionId=" + visit.visitId, 1)
            
            // Keep connection alive with periodic pings
            int pingCount = 0
            while (!response.isCommitted() && pingCount < 60) { // 5 minutes max
                Thread.sleep(5000) // Wait 5 seconds
                
                if (!response.isCommitted()) {
                    def pingData = [
                        type: "ping",
                        timestamp: System.currentTimeMillis(),
                        sessionId: visit.visitId,
                        architecture: "Visit-based sessions"
                    ]
                    sendSseEvent(response.writer, "ping", groovy.json.JsonOutput.toJson(pingData), pingCount + 2)
                    pingCount++
                }
            }
            
        } catch (InterruptedException e) {
            logger.info("SSE connection interrupted for session ${visit.visitId}")
            Thread.currentThread().interrupt()
        } catch (Exception e) {
            logger.warn("Enhanced SSE connection error: ${e.message}", e)
        } finally {
            // Clean up session - Visit persistence handles cleanup automatically
            try {
                def closeData = [
                    type: "disconnected", 
                    sessionId: visit.visitId,
                    timestamp: System.currentTimeMillis()
                ]
                sendSseEvent(response.writer, "disconnect", groovy.json.JsonOutput.toJson(closeData), -1)
            } catch (Exception e) {
                // Ignore errors during cleanup
            }
            
            // Complete async context if available
            if (request.isAsyncStarted()) {
                try {
                    request.getAsyncContext().complete()
                } catch (Exception e) {
                    logger.debug("Error completing async context: ${e.message}")
                }
            }
        }
    }
    
    private void handleMessage(HttpServletRequest request, HttpServletResponse response, ExecutionContextImpl ec) 
            throws IOException {
        
        // Get sessionId from request parameter or header
        String sessionId = request.getParameter("sessionId") ?: request.getHeader("Mcp-Session-Id")
        if (!sessionId) {
            response.setContentType("application/json")
            response.setCharacterEncoding("UTF-8")
            response.setStatus(HttpServletResponse.SC_BAD_REQUEST)
            response.writer.write(groovy.json.JsonOutput.toJson([
                error: "Missing sessionId parameter or header",
                architecture: "Visit-based sessions"
            ]))
            return
        }
        
        // Get Visit directly - this is our session
        def visit = ec.entity.find("moqui.server.Visit")
            .condition("visitId", sessionId)
            .one()
        
        if (!visit) {
            response.setContentType("application/json")
            response.setCharacterEncoding("UTF-8")
            response.setStatus(HttpServletResponse.SC_NOT_FOUND)
            response.writer.write(groovy.json.JsonOutput.toJson([
                error: "Session not found: " + sessionId,
                architecture: "Visit-based sessions"
            ]))
            return
        }
        
        // Verify user has access to this Visit
        if (visit.userId != ec.user.userId) {
            response.setContentType("application/json")
            response.setCharacterEncoding("UTF-8")
            response.setStatus(HttpServletResponse.SC_FORBIDDEN)
            response.writer.write(groovy.json.JsonOutput.toJson([
                error: "Access denied for session: " + sessionId,
                architecture: "Visit-based sessions"
            ]))
            return
        }
        
        // Create session wrapper for this Visit
        VisitBasedMcpSession session = new VisitBasedMcpSession(visit, response.writer, ec)
        
        try {
            // Read request body
            StringBuilder body = new StringBuilder()
            try {
                BufferedReader reader = request.getReader()
                String line
                while ((line = reader.readLine()) != null) {
                    body.append(line)
                }
            } catch (IOException e) {
                logger.error("Failed to read request body: ${e.message}")
                response.setContentType("application/json")
                response.setCharacterEncoding("UTF-8")
                response.setStatus(HttpServletResponse.SC_BAD_REQUEST)
                response.writer.write(groovy.json.JsonOutput.toJson([
                    jsonrpc: "2.0",
                    error: [code: -32700, message: "Failed to read request body: " + e.message],
                    id: null
                ]))
                return
            }
            
            String requestBody = body.toString()
            if (!requestBody.trim()) {
                response.setContentType("application/json")
                response.setCharacterEncoding("UTF-8")
                response.setStatus(HttpServletResponse.SC_BAD_REQUEST)
                response.writer.write(groovy.json.JsonOutput.toJson([
                    jsonrpc: "2.0",
                    error: [code: -32602, message: "Empty request body"],
                    id: null
                ]))
                return
            }
            
            // Parse JSON-RPC message
            def rpcRequest
            try {
                rpcRequest = jsonSlurper.parseText(requestBody)
            } catch (Exception e) {
                logger.error("Failed to parse JSON-RPC message: ${e.message}")
                response.setContentType("application/json")
                response.setCharacterEncoding("UTF-8")
                response.setStatus(HttpServletResponse.SC_BAD_REQUEST)
                response.writer.write(groovy.json.JsonOutput.toJson([
                    jsonrpc: "2.0",
                    error: [code: -32700, message: "Invalid JSON: " + e.message],
                    id: null
                ]))
                return
            }
            
            // Validate JSON-RPC 2.0 structure
            if (!rpcRequest?.jsonrpc || rpcRequest.jsonrpc != "2.0" || !rpcRequest?.method) {
                response.setContentType("application/json")
                response.setCharacterEncoding("UTF-8")
                response.setStatus(HttpServletResponse.SC_BAD_REQUEST)
                response.writer.write(groovy.json.JsonOutput.toJson([
                    jsonrpc: "2.0",
                    error: [code: -32600, message: "Invalid JSON-RPC 2.0 request"],
                    id: rpcRequest?.id ?: null
                ]))
                return
            }
            
            // Process method with session context
            def result = processMcpMethod(rpcRequest.method, rpcRequest.params, ec, sessionId)
            
            // Send response via MCP transport to the specific session
            def responseMessage = new JsonRpcResponse(result, rpcRequest.id)
            session.sendMessage(responseMessage)
            
            response.setContentType("application/json")
            response.setCharacterEncoding("UTF-8")
            response.setStatus(HttpServletResponse.SC_OK)
            response.writer.write(groovy.json.JsonOutput.toJson([
                jsonrpc: "2.0",
                id: rpcRequest.id,
                result: [status: "processed", sessionId: sessionId, architecture: "Visit-based"]
            ]))
            
        } catch (Exception e) {
            logger.error("Error processing message for session ${sessionId}: ${e.message}", e)
            response.setContentType("application/json")
            response.setCharacterEncoding("UTF-8")
            response.setStatus(HttpServletResponse.SC_INTERNAL_SERVER_ERROR)
            response.writer.write(groovy.json.JsonOutput.toJson([
                jsonrpc: "2.0",
                error: [code: -32603, message: "Internal error: " + e.message],
                id: null
            ]))
        }
    }
    
    private void handleJsonRpc(HttpServletRequest request, HttpServletResponse response, ExecutionContextImpl ec) 
            throws IOException {
        
        String method = request.getMethod()
        String acceptHeader = request.getHeader("Accept")
        String contentType = request.getContentType()
        
        logger.info("Enhanced MCP JSON-RPC Request: ${method} ${request.requestURI} - Accept: ${acceptHeader}, Content-Type: ${contentType}")
        
        // Handle POST requests for JSON-RPC
        if (!"POST".equals(method)) {
            response.setStatus(HttpServletResponse.SC_METHOD_NOT_ALLOWED)
            response.setContentType("application/json")
            response.writer.write(groovy.json.JsonOutput.toJson([
                jsonrpc: "2.0",
                error: [code: -32601, message: "Method Not Allowed. Use POST for JSON-RPC or GET /mcp-sse/sse for SSE."],
                id: null
            ]))
            return
        }
        
        // Read and parse JSON-RPC request
        String requestBody
        try {
            BufferedReader reader = request.getReader()
            StringBuilder body = new StringBuilder()
            String line
            while ((line = reader.readLine()) != null) {
                body.append(line)
            }
            requestBody = body.toString()
            
        } catch (IOException e) {
            logger.error("Failed to read request body: ${e.message}")
            response.setStatus(HttpServletResponse.SC_BAD_REQUEST)
            response.setContentType("application/json")
            response.writer.write(groovy.json.JsonOutput.toJson([
                jsonrpc: "2.0",
                error: [code: -32700, message: "Failed to read request body: " + e.message],
                id: null
            ]))
            return
        }
        
        if (!requestBody) {
            response.setStatus(HttpServletResponse.SC_BAD_REQUEST)
            response.setContentType("application/json")
            response.writer.write(groovy.json.JsonOutput.toJson([
                jsonrpc: "2.0",
                error: [code: -32602, message: "Empty request body"],
                id: null
            ]))
            return
        }
        
        def rpcRequest
        try {
            rpcRequest = jsonSlurper.parseText(requestBody)
        } catch (Exception e) {
            logger.error("Failed to parse JSON-RPC request: ${e.message}")
            response.setStatus(HttpServletResponse.SC_BAD_REQUEST)
            response.setContentType("application/json")
            response.writer.write(groovy.json.JsonOutput.toJson([
                jsonrpc: "2.0",
                error: [code: -32700, message: "Invalid JSON: " + e.message],
                id: null
            ]))
            return
        }
        
        // Validate JSON-RPC 2.0 structure
        if (!rpcRequest?.jsonrpc || rpcRequest.jsonrpc != "2.0" || !rpcRequest?.method) {
            response.setStatus(HttpServletResponse.SC_BAD_REQUEST)
            response.setContentType("application/json")
            response.writer.write(groovy.json.JsonOutput.toJson([
                jsonrpc: "2.0",
                error: [code: -32600, message: "Invalid JSON-RPC 2.0 request"],
                id: null
            ]))
            return
        }
        
        // Try to get session ID from cookie
        String sessionId = null
        def cookies = request.getCookies()
        if (cookies) {
            for (cookie in cookies) {
                if ("MCP-SESSION".equals(cookie.getName())) {
                    sessionId = cookie.getValue()
                    break
                }
            }
        }
        
        // Process MCP method using Moqui services with session ID if available
        def result = processMcpMethod(rpcRequest.method, rpcRequest.params, ec, sessionId)
        
        // Build JSON-RPC response
        def rpcResponse = [
            jsonrpc: "2.0",
            id: rpcRequest.id,
            result: result
        ]
        
        response.setContentType("application/json")
        response.setCharacterEncoding("UTF-8")
        
        // Set session cookie if result contains sessionId
        if (rpcResponse.result?.sessionId) {
            response.setHeader("Set-Cookie", "MCP-SESSION=${rpcResponse.result.sessionId}; Path=/; HttpOnly; SameSite=Lax")
        }
        
        response.writer.write(groovy.json.JsonOutput.toJson(rpcResponse))
    }
    
    private Map<String, Object> processMcpMethod(String method, Map params, ExecutionContextImpl ec, String sessionId) {
        logger.info("Enhanced METHOD: ${method} with params: ${params}, sessionId: ${sessionId}")
        
        try {
            // Ensure params is not null
            if (params == null) {
                params = [:]
            }
            
            // Add session context to parameters for services
            params.sessionId = sessionId
            
            switch (method) {
                case "initialize":
                    return callMcpService("mcp#Initialize", params, ec)
                case "ping":
                    return callMcpService("mcp#Ping", params, ec)
                case "tools/list":
                    return callMcpService("mcp#ToolsList", params, ec)
                case "tools/call":
                    return callMcpService("mcp#ToolsCall", params, ec)
                case "resources/list":
                    return callMcpService("mcp#ResourcesList", params, ec)
                case "resources/read":
                    return callMcpService("mcp#ResourcesRead", params, ec)
                case "notifications/initialized":
                    // Handle notification initialization - return success for now
                    return [initialized: true, sessionId: sessionId]
                case "notifications/send":
                    // Handle notification sending - return success for now  
                    return [sent: true, sessionId: sessionId]
                case "notifications/subscribe":
                    // Handle notification subscription - return success for now
                    return [subscribed: true, sessionId: sessionId]
                case "notifications/unsubscribe":
                    // Handle notification unsubscription - return success for now
                    return [unsubscribed: true, sessionId: sessionId]
                default:
                    throw new IllegalArgumentException("Method not found: ${method}")
            }
        } catch (Exception e) {
            logger.error("Error processing MCP method ${method}: ${e.message}", e)
            throw e
        }
    }
    
    private Map<String, Object> callMcpService(String serviceName, Map params, ExecutionContextImpl ec) {
        logger.info("Enhanced Calling MCP service: ${serviceName} with params: ${params}")
        
        try {
            def result = ec.service.sync().name("McpServices.${serviceName}")
                .parameters(params ?: [:])
                .call()
            
            logger.info("Enhanced MCP service ${serviceName} result: ${result}")
            if (result == null) {
                logger.error("Enhanced MCP service ${serviceName} returned null result")
                return [error: "Service returned null result"]
            }
            return result.result ?: [error: "Service result has no 'result' field"]
        } catch (Exception e) {
            logger.error("Error calling Enhanced MCP service ${serviceName}", e)
            return [error: e.message]
        }
    }
    
    private void sendSseEvent(PrintWriter writer, String eventType, String data, long eventId = -1) throws IOException {
        try {
            if (eventId >= 0) {
                writer.write("id: " + eventId + "\n")
            }
            writer.write("event: " + eventType + "\n")
            writer.write("data: " + data + "\n\n")
            writer.flush()
            
            if (writer.checkError()) {
                throw new IOException("Client disconnected")
            }
        } catch (Exception e) {
            throw new IOException("Failed to send SSE event: " + e.message, e)
        }
    }
    
    // CORS handling based on MoquiServlet pattern
    private static boolean handleCors(HttpServletRequest request, HttpServletResponse response, String webappName, ExecutionContextFactoryImpl ecfi) {
        String originHeader = request.getHeader("Origin")
        if (originHeader) {
            response.setHeader("Access-Control-Allow-Origin", originHeader)
            response.setHeader("Access-Control-Allow-Credentials", "true")
        }
        
        String methodHeader = request.getHeader("Access-Control-Request-Method")
        if (methodHeader) {
            response.setHeader("Access-Control-Allow-Methods", "GET, POST, OPTIONS")
            response.setHeader("Access-Control-Allow-Headers", "Content-Type, Authorization, Mcp-Session-Id, MCP-Protocol-Version, Accept")
            response.setHeader("Access-Control-Max-Age", "3600")
            return true
        }
        return false
    }
    
    @Override
    void destroy() {
        logger.info("Destroying EnhancedMcpServlet")
        
        // No session manager to shutdown - using Moqui's Visit system
        
        super.destroy()
    }
    
    /**
     * Broadcast message to all active sessions
     */
    void broadcastToAllSessions(JsonRpcMessage message) {
        // TODO: Implement broadcast using Moqui's Visit system if needed
        logger.info("Broadcast to all sessions not yet implemented")
    }
    
    /**
     * Get session statistics for monitoring
     */
    Map getSessionStatistics() {
        // TODO: Implement session statistics using Moqui's Visit system if needed
        return [activeSessions: 0, message: "Session statistics not yet implemented"]
    }
}