package com.gnx.telemedicine.util;

import jakarta.servlet.http.HttpServletRequest;
import jakarta.servlet.http.HttpServletResponse;
import org.slf4j.MDC;
import org.springframework.web.util.ContentCachingRequestWrapper;
import org.springframework.web.util.ContentCachingResponseWrapper;

import java.util.UUID;

/**
 * Utility class for managing correlation IDs across the application.
 * Uses MDC (Mapped Diagnostic Context) to store correlation IDs in the logging context.
 */
public class CorrelationIdUtil {
    private static final String CORRELATION_ID_HEADER = "X-Correlation-ID";
    private static final String CORRELATION_ID_MDC_KEY = "correlationId";
    private static final String REQUEST_ID_MDC_KEY = "requestId";
    private static final String USER_ID_MDC_KEY = "userId";
    private static final String REQUEST_PATH_MDC_KEY = "requestPath";
    private static final String REQUEST_METHOD_MDC_KEY = "requestMethod";

    /**
     * Generate a new correlation ID.
     */
    public static String generateCorrelationId() {
        return UUID.randomUUID().toString();
    }

    /**
     * Get correlation ID from request header or generate a new one.
     */
    public static String getOrGenerateCorrelationId(HttpServletRequest request) {
        String correlationId = request.getHeader(CORRELATION_ID_HEADER);
        if (correlationId == null || correlationId.isEmpty()) {
            correlationId = generateCorrelationId();
        }
        return correlationId;
    }

    /**
     * Set correlation ID in MDC and response header.
     */
    public static void setCorrelationId(String correlationId, HttpServletResponse response) {
        MDC.put(CORRELATION_ID_MDC_KEY, correlationId);
        response.setHeader(CORRELATION_ID_HEADER, correlationId);
    }

    /**
     * Set correlation ID in MDC.
     */
    public static void setCorrelationId(String correlationId) {
        MDC.put(CORRELATION_ID_MDC_KEY, correlationId);
    }

    /**
     * Get correlation ID from MDC.
     */
    public static String getCorrelationId() {
        return MDC.get(CORRELATION_ID_MDC_KEY);
    }

    /**
     * Set request ID in MDC.
     */
    public static void setRequestId(String requestId) {
        MDC.put(REQUEST_ID_MDC_KEY, requestId);
    }

    /**
     * Set user ID in MDC.
     */
    public static void setUserId(String userId) {
        if (userId != null && !userId.isEmpty()) {
            MDC.put(USER_ID_MDC_KEY, userId);
        }
    }

    /**
     * Set request path in MDC.
     */
    public static void setRequestPath(String path) {
        MDC.put(REQUEST_PATH_MDC_KEY, path);
    }

    /**
     * Set request method in MDC.
     */
    public static void setRequestMethod(String method) {
        MDC.put(REQUEST_METHOD_MDC_KEY, method);
    }

    /**
     * Set all request context in MDC.
     */
    public static void setRequestContext(HttpServletRequest request, String userId) {
        String correlationId = getOrGenerateCorrelationId(request);
        setCorrelationId(correlationId);
        setRequestPath(request.getRequestURI());
        setRequestMethod(request.getMethod());
        if (userId != null) {
            setUserId(userId);
        }
    }

    /**
     * Clear all MDC context.
     */
    public static void clear() {
        MDC.clear();
    }

    /**
     * Remove correlation ID from MDC.
     */
    public static void removeCorrelationId() {
        MDC.remove(CORRELATION_ID_MDC_KEY);
    }

    /**
     * Remove user ID from MDC.
     */
    public static void removeUserId() {
        MDC.remove(USER_ID_MDC_KEY);
    }
}

