/applinks-plugin/src/main/java/com/atlassian/applinks/core/auth/ApplicationLinkAnalyticsRequest.java

package com.atlassian.applinks.core.auth;

import com.atlassian.applinks.analytics.ApplinksRequestExecutionEvent;
import com.atlassian.applinks.api.ApplicationLink;
import com.atlassian.applinks.api.ApplicationLinkRequest;
import com.atlassian.applinks.api.ApplicationLinkResponseHandler;
import com.atlassian.event.api.EventPublisher;
import com.atlassian.sal.api.net.RequestFilePart;
import com.atlassian.sal.api.net.Response;
import com.atlassian.sal.api.net.ResponseException;
import com.atlassian.sal.api.net.ResponseHandler;
import com.atlassian.sal.api.net.ReturningResponseHandler;

import java.util.List;
import java.util.Map;
import java.util.Optional;

import static java.util.Arrays.stream;
import static java.util.Objects.requireNonNull;

/**
 * Wrapper for {@link ApplicationLinkRequest} to send analytics for each request.
 *
 * @since 5.4.24
 */
public class ApplicationLinkAnalyticsRequest implements ApplicationLinkRequest {
    private final ApplicationLink remoteApplink;
    private final EventPublisher publisher;

    private ApplicationLinkRequest wrappedRequest;
    private boolean unknownReqBodySize;
    private long reqBodySize;
    private int reqUrlSize;

    /**
     * Wraps a request to send analytics for each request.
     *
     * @param wrappedRequest Request to be wrapped.
     * @param remoteApplink The remoteApplink to the remote server.
     * @param publisher Event publisher for analytics.
     */
    ApplicationLinkAnalyticsRequest(final ApplicationLinkRequest wrappedRequest,
                                    final ApplicationLink remoteApplink,
                                    final EventPublisher publisher) {
        this.wrappedRequest = wrappedRequest;
        this.remoteApplink = requireNonNull(remoteApplink);
        this.publisher = requireNonNull(publisher);
        unknownReqBodySize = false;
        reqBodySize = 0;
        reqUrlSize = 0;
    }

    @Override
    public <R> R execute(final ApplicationLinkResponseHandler<R> responseHandler) throws ResponseException {
        ApplicationLinkAnalyticsResponseHandler<R> responseWrapper = new ApplicationLinkAnalyticsResponseHandler<>(responseHandler);
        R r = wrappedRequest.execute(responseWrapper);
        sendExecuteAnalytics(responseWrapper.getResponse()
                .map(ApplicationLinkAnalyticsRequest::getResponseSize)
                .orElse(null));
        return r;
    }

    @Override
    public ApplicationLinkRequest setConnectionTimeout(final int connectionTimeout) {
        wrappedRequest = wrappedRequest.setConnectionTimeout(connectionTimeout);
        return this;
    }

    @Override
    public ApplicationLinkRequest setSoTimeout(final int soTimeout) {
        wrappedRequest = wrappedRequest.setSoTimeout(soTimeout);
        return this;
    }

    @Override
    public ApplicationLinkRequest setUrl(final String url) {
        wrappedRequest = wrappedRequest.setUrl(url);
        this.reqUrlSize = url.getBytes().length;
        return this;
    }

    @Override
    public ApplicationLinkRequest setRequestBody(final String requestBody) {
        wrappedRequest = wrappedRequest.setRequestBody(requestBody);
        this.reqBodySize = requestBody.getBytes().length;
        this.unknownReqBodySize = false;
        return this;
    }

    @Override
    public ApplicationLinkRequest setRequestBody(final String requestBody, String contentType) {
        wrappedRequest = wrappedRequest.setRequestBody(requestBody, contentType);
        this.reqBodySize = requestBody.getBytes().length;
        this.unknownReqBodySize = false;
        return this;
    }

    @Override
    public ApplicationLinkRequest setFiles(final List<RequestFilePart> files) {
        wrappedRequest = wrappedRequest.setFiles(files);
        this.reqBodySize = files.stream()
                .map(f -> f.getFile().length())
                .mapToLong(Long::longValue)
                .sum();
        this.unknownReqBodySize = false;
        return this;
    }

    @Override
    public ApplicationLinkRequest setEntity(final Object entity) {
        wrappedRequest = wrappedRequest.setEntity(entity);
        // It's better not to return the size.
        this.reqBodySize = 0;
        this.unknownReqBodySize = true;
        return this;
    }

    @Override
    public ApplicationLinkRequest addRequestParameters(final String... params) {
        wrappedRequest = wrappedRequest.addRequestParameters(params);
        this.reqBodySize = stream(params)
                .map(p -> (long) p.getBytes().length)
                .mapToLong(Long::longValue)
                .sum();
        this.unknownReqBodySize = false;
        return this;
    }

    @Override
    public ApplicationLinkRequest addBasicAuthentication(final String hostname,
                                                         final String username,
                                                         final String password) {
        wrappedRequest = wrappedRequest.addBasicAuthentication(hostname, username, password);
        return this;
    }

    @Override
    public ApplicationLinkRequest addHeader(final String name, final String value) {
        wrappedRequest = wrappedRequest.addHeader(name, value);
        return this;
    }

    @Override
    public ApplicationLinkRequest setHeader(final String name, final String value) {
        wrappedRequest = wrappedRequest.setHeader(name, value);
        return this;
    }

    @Override
    public ApplicationLinkRequest setFollowRedirects(final boolean followRedirects) {
        wrappedRequest = wrappedRequest.setFollowRedirects(followRedirects);
        return this;
    }

    @Override
    public Map<String, List<String>> getHeaders() {
        return wrappedRequest.getHeaders();
    }

    @Override
    public void execute(final ResponseHandler<? super Response> responseHandler) throws ResponseException {
        AnalyticsResponseHandler<? super Response> responseWrapper
                = new AnalyticsResponseHandler<>(responseHandler);
        wrappedRequest.execute(responseWrapper);
        sendExecuteAnalytics(responseWrapper.getResponse()
                .map(ApplicationLinkAnalyticsRequest::getResponseSize)
                .orElse(null));
    }

    @Override
    public String execute() throws ResponseException {
        String result = wrappedRequest.execute();
        sendExecuteAnalytics(result != null ? (long) result.getBytes().length : null);
        return result;
    }

    @Override
    public <T> T executeAndReturn(final ReturningResponseHandler<? super Response, T> returningResponseHandler) throws ResponseException {
        ReturningAnalyticsResponseHandler<? super Response, T> responseWrapper
                = new ReturningAnalyticsResponseHandler<>(returningResponseHandler);
        T t = wrappedRequest.executeAndReturn(responseWrapper);
        sendExecuteAnalytics(responseWrapper.getResponse()
                .map(ApplicationLinkAnalyticsRequest::getResponseSize)
                .orElse(null));
        return t;
    }

    private void sendExecuteAnalytics(final Long approxResponseSize) {
        Long approxRequestSize;
        // If body size is unknown, send null, as payload size can't be estimated.
        if (unknownReqBodySize) {
            approxRequestSize = null;
        } else {
            approxRequestSize = reqBodySize + reqUrlSize + getHeaderSize();
        }

        publisher.publish(new ApplinksRequestExecutionEvent(approxRequestSize, approxResponseSize, remoteApplink.getId().get()));
    }

    private long getHeaderSize() {
        // Gets the size in bytes of the header by summing the length in bytes of
        // its various entries.
        return getHeaders().entrySet().stream()
                .map(entry -> {
                    long keyLength = entry.getKey().getBytes().length;
                    long valueLength = entry.getValue().stream()
                            .map(v -> (long) v.getBytes().length)
                            .mapToLong(Long::longValue)
                            .sum();
                    return keyLength + valueLength;
                }).mapToLong(Long::longValue).sum();
    }

    private static long getResponseSize(Response response) {
        long headerSize = response.getHeaders().entrySet().stream()
                .map(entry -> (long) entry.getKey().getBytes().length + entry.getValue().getBytes().length)
                .mapToLong(Long::longValue)
                .sum();
        long bodySize;
        try {
            bodySize = response.getResponseBodyAsString().getBytes().length;
        } catch (ResponseException e) {
            bodySize = 0;
        }

        return headerSize + bodySize;
    }

    protected static class ApplicationLinkAnalyticsResponseHandler<R> implements ApplicationLinkResponseHandler<R> {
        private final ApplicationLinkResponseHandler<R> wrappedResponseHandler;
        private Response response;

        ApplicationLinkAnalyticsResponseHandler(ApplicationLinkResponseHandler<R> wrappedResponseHandler) {
            this.wrappedResponseHandler = requireNonNull(wrappedResponseHandler);
            this.response = null;
        }

        public R credentialsRequired(Response response) throws ResponseException {
            return wrappedResponseHandler.credentialsRequired(response);
        }

        public R handle(Response response) throws ResponseException {
            this.response = response;
            return wrappedResponseHandler.handle(response);
        }

        Optional<Response> getResponse() {
            return Optional.ofNullable(response);
        }
    }

    protected static class AnalyticsResponseHandler<R extends Response> implements ResponseHandler<R> {
        private final ResponseHandler<R> wrappedResponseHandler;
        private R response;

        AnalyticsResponseHandler(ResponseHandler<R> wrappedResponseHandler) {
            this.wrappedResponseHandler = requireNonNull(wrappedResponseHandler);
            this.response = null;
        }

        @Override
        public void handle(R response) throws ResponseException {
            this.response = response;
            wrappedResponseHandler.handle(response);
        }

        Optional<Response> getResponse() {
            return Optional.ofNullable(response);
        }
    }

    protected static class ReturningAnalyticsResponseHandler<R extends Response, T> implements ReturningResponseHandler<R, T> {
        private final ReturningResponseHandler<R, T> wrappedResponseHandler;
        private R response;

        ReturningAnalyticsResponseHandler(ReturningResponseHandler<R, T> wrappedResponseHandler) {
            this.wrappedResponseHandler = requireNonNull(wrappedResponseHandler);
            this.response = null;
        }

        @Override
        public T handle(R response) throws ResponseException {
            this.response = response;
            return wrappedResponseHandler.handle(response);
        }

        Optional<Response> getResponse() {
            return Optional.ofNullable(response);
        }
    }
}