GlobalResponseWrapper.java

package io.github.og4dev.advice;

import io.github.og4dev.annotation.AutoResponse;
import io.github.og4dev.dto.ApiResponse;
import org.jspecify.annotations.NullMarked;
import org.jspecify.annotations.Nullable;
import org.springframework.core.MethodParameter;
import org.springframework.http.HttpStatus;
import org.springframework.http.MediaType;
import org.springframework.http.ProblemDetail;
import org.springframework.http.ResponseEntity;
import org.springframework.http.converter.HttpMessageConverter;
import org.springframework.http.server.ServerHttpRequest;
import org.springframework.http.server.ServerHttpResponse;
import org.springframework.http.server.ServletServerHttpResponse;
import org.springframework.web.bind.annotation.RestControllerAdvice;
import org.springframework.web.servlet.mvc.method.annotation.ResponseBodyAdvice;
import tools.jackson.core.JacksonException;
import tools.jackson.databind.ObjectMapper;

/**
 * Global response interceptor that automatically wraps REST controller outputs into the
 * standardized {@link ApiResponse} format.
 * <p>
 * This wrapper is conditionally activated <b>only</b> for controllers or specific methods
 * annotated with the {@link AutoResponse @AutoResponse} annotation. It provides a seamless
 * developer experience by eliminating the need to manually return {@code ResponseEntity<ApiResponse<T>>}
 * from every controller method.
 * </p>
 * <h2>Core Functionalities:</h2>
 * <ul>
 * <li><b>Automatic Encapsulation:</b> Intercepts raw DTOs, Lists, or primitive responses and packages
 * them into the {@code content} field of an {@code ApiResponse}.</li>
 * <li><b>Status Code Preservation:</b> Dynamically reads the current HTTP status of the response
 * (e.g., set via {@code @ResponseStatus(HttpStatus.CREATED)}) and ensures it is accurately
 * reflected in the final {@code ApiResponse}.</li>
 * <li><b>String Payload Compatibility:</b> Safely intercepts raw {@code String} returns and manually
 * serializes them to prevent {@code ClassCastException} when Spring utilizes the {@code StringHttpMessageConverter}.</li>
 * <li><b>Safety Mechanisms:</b> Intelligently skips wrapping if the response is already formatted
 * to prevent double-wrapping errors or interference with standard error handling protocols.</li>
 * </ul>
 *
 * @author Pasindu OG
 * @version 1.4.0
 * @see AutoResponse
 * @see ApiResponse
 * @see ResponseBodyAdvice
 * @since 1.4.0
 */
@RestControllerAdvice
@SuppressWarnings("unused")
public @NullMarked class GlobalResponseWrapper implements ResponseBodyAdvice<Object> {

    private final ObjectMapper objectMapper;

    /**
     * Constructs a new {@code GlobalResponseWrapper} with the provided {@link ObjectMapper}.
     *
     * @param objectMapper The Jackson object mapper used for explicit string serialization.
     */
    public GlobalResponseWrapper(ObjectMapper objectMapper) {
        this.objectMapper = objectMapper;
    }

    /**
     * Determines whether the current response should be intercepted and wrapped.
     * <p>
     * This method evaluates two main conditions before allowing the response to be wrapped:
     * </p>
     * <ol>
     * <li><b>Annotation Presence:</b> The target controller class or the specific handler method
     * must be annotated with {@link AutoResponse}.</li>
     * <li><b>Type Exclusion:</b> The return type must <b>not</b> be one of the explicitly excluded types.</li>
     * </ol>
     * <p>
     * To guarantee application stability and adherence to standard HTTP protocols, this method
     * specifically <b>excludes</b> the following return types from being wrapped:
     * </p>
     * <ul>
     * <li>{@link ApiResponse} - Prevents recursive double-wrapping (e.g., {@code ApiResponse<ApiResponse<T>>}).</li>
     * <li>{@link ResponseEntity} - Skips manual responses to respect developer's explicit configurations.</li>
     * <li>{@link ProblemDetail} - Excludes RFC 9457 error responses generated by exception handlers.</li>
     * </ul>
     * <p>
     * Note: Unlike standard wrappers, raw {@link String} payloads are <b>supported</b> and handled
     * appropriately during the write phase.
     * </p>
     *
     * @param returnType    The return type of the controller method.
     * @param converterType The selected HTTP message converter.
     * @return {@code true} if annotated with {@code @AutoResponse} and not an excluded type; {@code false} otherwise.
     */
    @Override
    public @NullMarked boolean supports(MethodParameter returnType, Class<? extends HttpMessageConverter<?>> converterType) {
        Class<?> type = returnType.getParameterType();
        boolean isExcludedType = ApiResponse.class.isAssignableFrom(type) ||
                ResponseEntity.class.isAssignableFrom(type) ||
                ProblemDetail.class.isAssignableFrom(type);
        if (isExcludedType) return false;
        boolean hasClassAnnotation = returnType.getDeclaringClass().isAnnotationPresent(AutoResponse.class);
        boolean hasMethodAnnotation = returnType.hasMethodAnnotation(AutoResponse.class);
        return hasClassAnnotation || hasMethodAnnotation;
    }

    /**
     * Intercepts the response body before it is written to the output stream and encapsulates it
     * within an {@link ApiResponse}.
     * <p>
     * This method extracts the actual HTTP status code set on the current response (defaulting to 200 OK).
     * Based on whether the status code represents a success (2xx) or another state, it dynamically
     * assigns an appropriate message ("Success" or "Processed") to the API response.
     * </p>
     * <p>
     * <b>Special String Handling:</b> If the intercepted payload is a raw {@code String}, it is
     * explicitly serialized to a JSON string using the configured {@link ObjectMapper}, and the
     * response {@code Content-Type} is strictly set to {@code application/json}. This prevents
     * standard message converter conflicts.
     * </p>
     *
     * @param body                  The raw object returned by the controller method.
     * @param returnType            The return type of the controller method.
     * @param selectedContentType   The selected content type for the response.
     * @param selectedConverterType The selected HTTP message converter.
     * @param request               The current server HTTP request.
     * @param response              The current server HTTP response.
     * @return The newly wrapped {@code ApiResponse} object ready to be serialized, or a pre-serialized
     * JSON {@code String} if the original payload was a raw string.
     */
    @Override
    public @Nullable Object beforeBodyWrite(@Nullable Object body, MethodParameter returnType, MediaType selectedContentType, Class<? extends HttpMessageConverter<?>> selectedConverterType, ServerHttpRequest request, ServerHttpResponse response) {
        int statusCode = HttpStatus.OK.value();

        if (response instanceof ServletServerHttpResponse serverHttpResponse) {
            statusCode = serverHttpResponse.getServletResponse().getStatus();
        }
        HttpStatus httpStatus = HttpStatus.valueOf(statusCode);
        ApiResponse<Object> apiResponse = ApiResponse.status(httpStatus.is2xxSuccessful() ? "Success" : "Processed", body, httpStatus).getBody();

        if (body instanceof String) {
            try {
                response.getHeaders().setContentType(MediaType.APPLICATION_JSON);
                assert apiResponse != null;
                return objectMapper.writeValueAsString(apiResponse);
            } catch (JacksonException e) {
                return body;
            }
        }
        return apiResponse;
    }
}