package com.atlassian.stash.internal.web.filters;

import com.atlassian.bitbucket.request.RequestManager;
import com.atlassian.stash.internal.web.HttpRequestInfoHelper;
import com.atlassian.stash.internal.web.util.web.CountingHttpServletRequestWrapper;
import com.atlassian.stash.internal.web.util.web.CountingHttpServletResponseWrapper;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;
import org.springframework.web.filter.OncePerRequestFilter;

import javax.servlet.FilterChain;
import javax.servlet.ServletException;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.io.IOException;

/**
 * Adds attributes about the current {@code HttpServletRequest} to the logging MDC and as response headers. Be careful
 * when moving this filter in the filter chain, because this operation needs to happen as close to the boundaries of
 * the request as possible.
 * <p>
 * Note: This consolidates MDC and headers in a single filter purely to try and reduce the overall depth of the filter
 * chain. Given that the values and concepts are all the same, just with different targets, it seems wasteful to use
 * multiple filters to do the work.
 */
@Component("requestAttributeFilter")
public class RequestAttributeFilter extends OncePerRequestFilter {

    public static final String X_AREQUESTID = "X-AREQUESTID";
    public static final String X_ASESSIONID = "X-ASESSIONID";

    private final HttpRequestInfoHelper helper;
    private final RequestManager requestManager;

    @Autowired
    public RequestAttributeFilter(HttpRequestInfoHelper helper, RequestManager requestManager) {
        this.helper = helper;
        this.requestManager = requestManager;
    }

    @Override
    protected void doFilterInternal(final HttpServletRequest request, final HttpServletResponse response,
                                    final FilterChain filterChain) throws ServletException, IOException {
        try {
            requestManager.doAsRequest(requestContext -> {
                //Set response headers for the request ID, and the session ID if one is available
                response.setHeader(X_AREQUESTID, requestContext.getId());
                if (requestContext.hasSessionId()) {
                    response.setHeader(X_ASESSIONID, requestContext.getSessionId());
                }

                CountingHttpServletRequestWrapper requestWrapper = new CountingHttpServletRequestWrapper(request);
                CountingHttpServletResponseWrapper responseWrapper = new CountingHttpServletResponseWrapper(response);

                filterChain.doFilter(requestWrapper, responseWrapper);

                //Record the number of bytes sent/received and http response code in the access log
                responseWrapper.flushBuffer();
                requestContext.setBytesRead(requestWrapper.getBytesRead());
                requestContext.setBytesWritten(responseWrapper.getBytesWritten());
                requestContext.setResponseCode(responseWrapper.getStatus());

                return null;
            }, helper.createRequestInfoProvider(request, response));
        } catch (Exception e) {
            if (e instanceof ServletException) {
                throw (ServletException) e;
            }
            if (e instanceof IOException) {
                throw (IOException) e;
            }
            throw (RuntimeException) e;
        }
    }

    /**
     * In order to render /mvc/error404 and /mvc/error500 in an error dispatch
     * {@link RequestManager#getRequestContext()}
     * must not be null. Therefore we must ensure error dispatches are filtered
     */
    @Override
    protected boolean shouldNotFilterErrorDispatch() {
        return false;
    }
}
