package com.x.x.x;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.slf4j.MDC;
import org.springframework.core.annotation.Order;
import org.springframework.http.HttpHeaders;
import org.springframework.stereotype.Component;
import org.springframework.web.util.ContentCachingResponseWrapper;
import javax.servlet.*;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.io.IOException;
import java.util.Collections;
import java.util.UUID;
import java.util.function.Function;
import java.util.stream.Collectors;

/**
 * @author vishal.mamidi
 *
 */
@Order(1)
@Component
public class LoggerFilter implements Filter {

    private final Logger LOGGER = LoggerFactory.getLogger(LoggerFilter.class);
    private final String X_REQUEST_ID = "X-Request-ID";

    @Override
    public void doFilter(ServletRequest request,
                         ServletResponse response,
                         FilterChain chain) throws IOException, ServletException {
        HttpServletRequest req = (HttpServletRequest) request;
        HttpServletResponse res = (HttpServletResponse) response;
        String responseStr = null;
        try {
            if (!"GET".equalsIgnoreCase(req.getMethod()))
                MDC.put("RequestBody", request.getReader().lines().collect(Collectors.joining(System.lineSeparator())));
            HttpHeaders httpHeaders = Collections
                    .list(req.getHeaderNames())
                    .stream()
                    .collect(Collectors.toMap(
                            Function.identity(),
                            h -> Collections.list(req.getHeaders(h)),
                            (oldValue, newValue) -> newValue,
                            HttpHeaders::new
                    ));

            MDC.put("Headers", String.valueOf(httpHeaders));
            addXRequestId(req);
            LOGGER.info("Type: Request - path: {}, method: {}, query: {}, host: {}",
                    req.getRequestURI(), req.getMethod(), req.getQueryString(), req.getRemoteHost());
            res.setHeader(X_REQUEST_ID, MDC.get(X_REQUEST_ID));

            ContentCachingResponseWrapper responseCacheWrapperObject = new ContentCachingResponseWrapper((HttpServletResponse) response);

            chain.doFilter(request, responseCacheWrapperObject);
            byte[] responseArray = responseCacheWrapperObject.getContentAsByteArray();
            responseStr = new String(responseArray, responseCacheWrapperObject.getCharacterEncoding());
            responseCacheWrapperObject.copyBodyToResponse();

        } finally {
            MDC.remove("Headers");
            MDC.remove("RequestBody");
            MDC.put("ResponseBody", responseStr);
            HttpHeaders httpHeaders = Collections
                    .synchronizedCollection(res.getHeaderNames())
                    .stream()
                    .collect(Collectors.toMap(
                            Function.identity(),
                            h -> Collections.list(req.getHeaders(h)),
                            (oldValue, newValue) -> newValue,
                            HttpHeaders::new
                    ));
            MDC.put("Headers", String.valueOf(httpHeaders));
            LOGGER.info("Type: Response - statusCode: {}, path: {}, method: {}, query: {}",
                    res.getStatus(), req.getRequestURI(), req.getMethod(), req.getQueryString());
            MDC.clear();
        }
    }

    private void addXRequestId(HttpServletRequest request) {
        String xRequestId = request.getHeader(X_REQUEST_ID);
        if (xRequestId == null) {
            MDC.put(X_REQUEST_ID, UUID.randomUUID().toString().replace("-",""));
        } else {
            MDC.put(X_REQUEST_ID, xRequestId);
        }
    }
}