package com.atlassian.bitbucket.mesh.request;

import com.atlassian.bitbucket.mesh.execution.ExecutionContext;
import com.atlassian.bitbucket.mesh.execution.ExecutionLocal;
import com.atlassian.bitbucket.mesh.execution.ExecutionManager;
import com.atlassian.bitbucket.mesh.git.exception.FailedPreconditionException;
import com.atlassian.bitbucket.mesh.grpc.MetadataProvider;
import com.atlassian.bitbucket.mesh.logback.AccessLog;
import com.atlassian.bitbucket.mesh.logback.LogbackUtils;
import io.grpc.Metadata;

import javax.annotation.Nonnull;
import javax.annotation.Nullable;
import java.time.Duration;
import java.util.Collections;
import java.util.Optional;
import java.util.SortedSet;
import java.util.TreeSet;

import static com.atlassian.bitbucket.mesh.rpc.MetadataKeys.KEY_REMOTE_REQUEST;
import static com.atlassian.bitbucket.mesh.rpc.MetadataKeys.KEY_REMOTE_USER;
import static java.util.Objects.requireNonNull;

/**
 * Uses the {@link ExecutionManager} and {@link ExecutionContext} to track {@link RequestContext requests}.
 * <p>
 * This manager handles some additional request-related processing:
 * <ul>
 *     <li>Binds itself to the {@link LogbackUtils#bind logging context} so request details can be included
 *     in logs</li>
 *     <li>Uses {@link AccessLog} to write start and end lines to the access log</li>
 * </ul>
 */
public class DefaultRequestManager implements RequestManager, MetadataProvider {

    private final ExecutionLocal<RequestContext> contextLocal;
    private final ExecutionManager executionManager;

    public DefaultRequestManager(ExecutionManager executionManager) {
        this.executionManager = executionManager;

        contextLocal = executionManager.newExecutionLocal();

        //Bind the RequestManager to the logging context so it can be used to add request details to logs
        LogbackUtils.bind(this);
    }

    @Override
    public void apply(@Nonnull Metadata metadata) {
        requireNonNull(metadata, "metadata");

        contextLocal.ifPresent(context -> {
            String correlationId = context.getCorrelationId();
            String user = context.getUser();

            //If we have a correlationId, prepend that to our local request ID so we send both
            StringBuilder builder = new StringBuilder();
            if (correlationId != null) {
                builder.append(correlationId).append(',');
            }
            builder.append(context.getId());

            metadata.put(KEY_REMOTE_REQUEST, builder.toString());
            if (user != null) {
                metadata.put(KEY_REMOTE_USER, user);
            }
        });
    }

    @Nonnull
    @Override
    public Optional<RequestContext> getContext() {
        return Optional.ofNullable(contextLocal.get());
    }

    @Nonnull
    @Override
    public RequestContext getContextOrThrow() {
        return getContext().orElseThrow(() -> new FailedPreconditionException("No RequestContext is currently active"));
    }

    @Override
    public <T> T withContext(@Nonnull RequestDescriptor descriptor, @Nonnull RequestCallback<T> callback) {
        requireNonNull(descriptor, "descriptor");
        requireNonNull(callback, "callback");

        ExecutionContext executionContext = executionManager.getContext().orElse(null);
        if (executionContext == null) {
            //If there's no current ExecutionContext, create one. It will automatically be torn down after
            //the RequestCallback is invoked
            return executionManager.withContext(ctx -> callWithRequest(ctx, descriptor, callback));
        }

        //Otherwise, attach the request to the current ExecutionContext. Note that, for this path, the
        //RequestContext will _intentionally_ still be active after this method returns
        return callWithRequest(executionContext, descriptor, callback);
    }

    private <T> T callWithRequest(ExecutionContext executionContext, RequestDescriptor descriptor,
                                  RequestCallback<T> callback) {
        RequestContext requestContext = new DelegatingRequestContext(descriptor, executionContext);
        contextLocal.set(requestContext);

        AccessLog.start(requestContext);
        executionContext.addCleanupCallback(() -> AccessLog.end(requestContext));

        return callback.callWithRequest(requestContext);
    }

    private static class DelegatingRequestContext implements RequestContext {

        private final RequestDescriptor delegate;
        private final ExecutionContext executionContext;
        private final SortedSet<String> labels;
        private final long start;

        DelegatingRequestContext(RequestDescriptor delegate, ExecutionContext executionContext) {
            this.delegate = delegate;
            this.executionContext = executionContext;

            labels = new TreeSet<>();
            start = System.currentTimeMillis();
        }

        @Override
        public void addLabels(@Nonnull String first, @Nullable String... others) {
            labels.add(requireNonNull(first, "label"));

            if (others != null && others.length > 0) {
                for (int i = 0; i < others.length; ++i) {
                    labels.add(requireNonNull(others[i], "labels[" + i + "]"));
                }
            }
        }

        @Nonnull
        @Override
        public String getAction() {
            return delegate.getAction();
        }

        @Override
        public String getAgent() {
            return delegate.getAgent();
        }

        @Override
        public long getBytesReceived() {
            return delegate.getBytesReceived();
        }

        @Override
        public long getBytesSent() {
            return delegate.getBytesSent();
        }

        @Override
        public String getCorrelationId() {
            return delegate.getCorrelationId();
        }

        @Nonnull
        @Override
        public Duration getDuration() {
            return Duration.ofMillis(System.currentTimeMillis() - start);
        }

        @Nonnull
        @Override
        public String getId() {
            return executionContext.getId();
        }

        @Nonnull
        @Override
        public SortedSet<String> getLabels() {
            return Collections.unmodifiableSortedSet(labels);
        }

        @Override
        public int getMessagesReceived() {
            return delegate.getMessagesReceived();
        }

        @Override
        public int getMessagesSent() {
            return delegate.getMessagesSent();
        }

        @Nonnull
        @Override
        public String getProtocol() {
            return delegate.getProtocol();
        }

        @Nonnull
        @Override
        public Object getRawRequest() {
            return delegate.getRawRequest();
        }

        @Override
        public String getRemoteAddress() {
            return delegate.getRemoteAddress();
        }

        @Override
        public Integer getResult() {
            return delegate.getResult();
        }

        @Override
        public String getUser() {
            return delegate.getUser();
        }
    }
}
