package com.atlassian.bitbucket.mesh.git;

import com.atlassian.bitbucket.mesh.event.EventPublisher;
import com.atlassian.bitbucket.mesh.git.builder.GitQuarantineEnvironment;
import com.atlassian.bitbucket.mesh.git.exception.InvalidArgumentException;
import com.atlassian.bitbucket.mesh.git.hook.*;
import com.atlassian.bitbucket.mesh.git.hook.handler.AbstractGitHookHandler;
import com.atlassian.bitbucket.mesh.git.hosting.FilteringGitReceivePackStream;
import com.atlassian.bitbucket.mesh.git.hosting.GitReceivePackFilter;
import com.atlassian.bitbucket.mesh.git.hosting.GitReceivePackListener;
import com.atlassian.bitbucket.mesh.git.hosting.ReceivePackRequestAnalyzer;
import com.atlassian.bitbucket.mesh.git.ref.RefUpdate;
import com.atlassian.bitbucket.mesh.management.SidebandManager;
import com.atlassian.bitbucket.mesh.process.OsProcess;
import com.atlassian.bitbucket.mesh.repository.RefsUpdatedEvent;
import com.atlassian.bitbucket.mesh.repository.Repository;
import com.atlassian.bitbucket.mesh.request.RequestContext;
import com.atlassian.bitbucket.mesh.request.RequestManager;
import com.atlassian.bitbucket.mesh.rpc.ZeroCopyUtils;
import com.atlassian.bitbucket.mesh.rpc.v1.*;
import com.atlassian.bitbucket.mesh.transaction.UpdateRefsContext;
import com.atlassian.bitbucket.mesh.transaction.UpdateRefsContextHolder;
import com.atlassian.bitbucket.mesh.util.SatiableConsumer;
import com.google.common.collect.ImmutableList;
import com.google.protobuf.ByteString;
import io.grpc.stub.StreamObserver;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.annotation.Nonnull;
import java.io.IOException;
import java.nio.ByteBuffer;
import java.time.Duration;
import java.util.*;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;

import static com.atlassian.bitbucket.mesh.git.hook.GitHook.POST_RECEIVE;
import static com.atlassian.bitbucket.mesh.git.hook.GitHook.PRE_RECEIVE;
import static java.util.Objects.requireNonNull;
import static java.util.stream.Collectors.toList;

public abstract class AbstractHostingWriteFragmentHandler<ReqT>
        extends AbstractHostingFragmentHandler<ReqT, RpcHostingWriteStreamFragment> {

    private static final Logger log = LoggerFactory.getLogger(AbstractHostingWriteFragmentHandler.class);

    private final long maxFragmentSize;
    private final MutableRpcStreamChunk nextChunk;
    private final PostReceiveHookHandler postReceive;
    private final PreReceiveHookHandler preReceive;
    private final GrpcPushTracer pushTracer;
    private final ReceivePackRequestAnalyzer receivePackAnalyzer;
    private final List<GitReceivePackListener> receivePackListeners;
    private final RequestManager requestManager;
    private final SidebandManager sidebandManager;
    private final UpdateRefsContextHolder updateRefsContextHolder;

    private FilteringGitReceivePackStream filteringReceivePackStream;
    private Repository repository;

    protected AbstractHostingWriteFragmentHandler(@Nonnull RpcHostingOperationType defaultOperationType,
                                                  @Nonnull EventPublisher eventPublisher,
                                                  @Nonnull GitHookService hookService, @Nonnull Duration hookTimeout,
                                                  long maxFragmentSize,
                                                  @Nonnull StreamObserver<RpcHostingWriteStreamFragment> observer,
                                                  @Nonnull GrpcPushTracer pushTracer,
                                                  @Nonnull GitQuarantineRegistry quarantineRegistry,
                                                  @Nonnull List<GitReceivePackListener> receivePackListeners,
                                                  @Nonnull RequestManager requestManager,
                                                  @Nonnull SidebandManager sidebandManager,
                                                  @Nonnull UpdateRefsContextHolder updateRefsContextHolder) {
        super(defaultOperationType, hookService);

        requireNonNull(hookTimeout, "hookTimeout");
        requireNonNull(quarantineRegistry, "quarantineRegistry");

        this.maxFragmentSize = maxFragmentSize;
        this.pushTracer = requireNonNull(pushTracer, "pushTracer");
        this.receivePackListeners = requireNonNull(receivePackListeners, "receivePackListeners");
        this.requestManager = requireNonNull(requestManager, "requestManager");
        this.sidebandManager = sidebandManager;
        this.updateRefsContextHolder = requireNonNull(updateRefsContextHolder, "updateRefsContextHolder");

        nextChunk = new MutableRpcStreamChunk();

        postReceive = new PostReceiveHookHandler(
                maxFragmentSize, hookTimeout, eventPublisher, observer, this.updateRefsContextHolder);
        preReceive = new PreReceiveHookHandler(
                quarantineRegistry, updateRefsContextHolder, maxFragmentSize, hookTimeout, observer);
        receivePackAnalyzer = new ReceivePackRequestAnalyzer();
    }

    @Nonnull
    @Override
    public RpcHostingWriteStreamFragment createResultFragment(int exitCode) {
        RpcHostingStreamResult result = getResult(exitCode);
        pushTracer.onResult(result);
        return RpcHostingWriteStreamFragment.newBuilder()
                .setResult(result)
                .build();
    }

    @Nonnull
    @Override
    public RpcHostingWriteStreamFragment createStartFragment(@Nonnull OsProcess process) {
        return RpcHostingWriteStreamFragment.newBuilder()
                .setStart(RpcStreamStart.newBuilder()
                        .setCommandLine(ByteString.copyFromUtf8(process.toString())))
                .build();
    }

    @Nonnull
    @Override
    public RpcHostingWriteStreamFragment createStderrFragment(@Nonnull RpcStreamChunk chunk) {
        pushTracer.onStderr(chunk);
        return RpcHostingWriteStreamFragment.newBuilder()
                .setStderr(chunk)
                .build();
    }

    @Nonnull
    @Override
    public RpcHostingWriteStreamFragment createStdoutFragment(@Nonnull RpcStreamChunk chunk) {
        pushTracer.onStdout(chunk);
        return RpcHostingWriteStreamFragment.newBuilder()
                .setStdout(chunk)
                .build();
    }

    @Override
    public void onComplete() {
        wrapUp();
    }

    @Override
    public void onError(boolean canceled, Throwable t) {
        wrapUp();
    }

    protected void invokeReceivePackListeners(Repository repository) {
        List<GitReceivePackFilter> filters = receivePackListeners.stream()
                .map(listener -> listener.createFilter(repository))
                .filter(Objects::nonNull)
                .collect(toList());
        filteringReceivePackStream = new FilteringGitReceivePackStream(nextChunk);
        if (!filters.isEmpty()) {
            filters.forEach(filteringReceivePackStream::addFilter);
        }
        //Add the analyzer last so it processes the transformed stdin
        filteringReceivePackStream.addFilter(receivePackAnalyzer);
    }

    @Override
    protected boolean isUpdate() {
        return receivePackAnalyzer.isUpdate();
    }

    protected void onPostReceive(@Nonnull RpcHostingPostReceiveResponse resp) {
        pushTracer.onPostReceive(0);
        postReceive.complete(requireNonNull(ZeroCopyUtils.deepCopy(resp), "response"));
    }

    protected void onPreReceive(@Nonnull RpcHostingPreReceiveResponse resp) {
        pushTracer.onPreReceive(resp.getRejected() ? 1 : 0);
        preReceive.complete(requireNonNull(ZeroCopyUtils.deepCopy(resp), "response"));
    }

    protected RpcStreamChunk onStdin(@Nonnull RpcStreamChunk chunk) {
        if (filteringReceivePackStream == null) {
            return chunk;
        }
        filteringReceivePackStream.write(chunk);
        return nextChunk.getAndReset();
    }

    @Nonnull
    @Override
    protected HookRequestHandle registerHookRequest(Repository repository) {
        this.repository = repository;
        return registerHookRequest(repository, preReceive, postReceive);
    }

    private static List<List<RpcRefUpdate>> parseRefUpdates(HookStdin stdin, long maxFragmentSize) {
        try {
            return stdin.readLines(new RefChangeLineProcessor(new RefChangeFragmentChunker(maxFragmentSize)));
        } catch (IOException e) {
            throw new IllegalStateException("Failed to read and parse ref updates", e);
        } finally {
            // Whether we parse the updates successfully or not, we should never need access to
            // stdin again after this point
            stdin.release();
        }
    }

    private void wrapUp() {
        UpdateRefsContext context = updateRefsContextHolder.get();
        try {
            if (context == null) {
                return;
            }
            Set<String> committedRefUpdates = context.getCommittedRefUpdates();
            Optional<RequestContext> maybeContext = requestManager.getContext();
            if (!maybeContext.isPresent()) {
                log.warn("Pending ref updates cannot be sent back to the control plane due to missing " +
                        "request context.");
                return;
            }
            Map<String, RefUpdate> updatedRefs = context.getRefUpdates();
            if (!committedRefUpdates.isEmpty()) {
                RefChangeFragmentChunker chunker = new RefChangeFragmentChunker(maxFragmentSize);
                committedRefUpdates.forEach(refId -> {
                    RefUpdate refUpdate = updatedRefs.get(refId);
                    if (refUpdate == null) {
                        log.debug("Ref '{}' was committed but was not included in the pre-receive hook.", refId);
                        return;
                    }
                    chunker.addRefUpdate(refId, refUpdate.getOldHash(), refUpdate.getNewHash());
                });
                Iterator<List<RpcRefUpdate>> it = chunker.getResult().iterator();
                while (it.hasNext()) {
                    List<RpcRefUpdate> refUpdates = it.next();
                    sidebandManager.send(RpcSidebandFollowerFragment.newBuilder()
                            .setOutOfBandRefUpdates(RpcOutOfBandRefUpdates.newBuilder()
                                    .setCorrelationId(maybeContext.get().getCorrelationId())
                                    .setUser(maybeContext.get().getUser())
                                    .setRepository(repository.getId())
                                    .addAllRefUpdates(refUpdates)
                                    .setMoreRefUpdates(it.hasNext()))
                            .build());
                }
            }
        } finally {
            updateRefsContextHolder.clear();
        }
    }

    private abstract static class CompletableGitHookHandler<T> extends AbstractGitHookHandler {

        private static final Logger log = LoggerFactory.getLogger(CompletableGitHookHandler.class);
        private final CompletableFuture<T> future;
        private final Duration timeout;

        protected CompletableGitHookHandler(@Nonnull GitHook hook, @Nonnull Duration timeout) {
            super(hook);

            this.timeout = requireNonNull(timeout, "timeout");
            future = new CompletableFuture<>();
        }

        public void complete(@Nonnull T value) {
            future.complete(requireNonNull(value, "value"));
        }

        @Override
        public boolean handle(@Nonnull Repository repository, @Nonnull HookRequest request,
                              @Nonnull GitHookOutputWriter writer) {
            try {
                T response = future.get(timeout.toMillis(), TimeUnit.MILLISECONDS);
                List<RpcHookChunk> output = readOutput(response);

                for (RpcHookChunk chunk : output) {
                    switch (chunk.getChannel()) {
                        case CHANNEL_STDERR:
                            for (ByteBuffer buffer : chunk.getData().asReadOnlyByteBufferList()) {
                                writer.writeStderr(buffer);
                            }
                            break;
                        case CHANNEL_STDOUT:
                            for (ByteBuffer buffer : chunk.getData().asReadOnlyByteBufferList()) {
                                writer.writeStdout(buffer);
                            }
                            break;
                        default:
                            throw new InvalidArgumentException("Received [" + chunk.getChannel() + "] where " +
                                    "[case CHANNEL_STDERR or CHANNEL_STDOUT] was expected");
                    }
                }

                return isAccepted(response);
            } catch (IOException e) {
                log.error("Failed to process git hooks", e);
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt();
            } catch (ExecutionException executionException) {
                log.error("Git hook execution resulted in exception", executionException.getCause());
            } catch (TimeoutException e) {
                log.error("Timeout waiting for response from RPC client for git hook", e);
            }

            return false;
        }

        protected abstract boolean isAccepted(@Nonnull T response);

        @Nonnull
        protected abstract List<RpcHookChunk> readOutput(@Nonnull T response);
    }

    private static class MutableRpcStreamChunk implements SatiableConsumer<RpcStreamChunk> {

        private RpcStreamChunk value;

        @Override
        public boolean accept(@Nonnull RpcStreamChunk chunk) {
            value = value == null ? chunk : chunk.toBuilder()
                    .setData(value.getData().concat(chunk.getData()))
                    .build();
            return false;
        }

        RpcStreamChunk getAndReset() {
            RpcStreamChunk result = value;
            value = null;
            return result;
        }
    }

    private static class PostReceiveHookHandler extends CompletableGitHookHandler<RpcHostingPostReceiveResponse> {

        private final EventPublisher eventPublisher;
        private final long maxFragmentSize;
        private final StreamObserver<RpcHostingWriteStreamFragment> observer;
        private final UpdateRefsContextHolder updateRefsContextHolder;

        PostReceiveHookHandler(long maxFragmentSize, @Nonnull Duration timeout, @Nonnull EventPublisher eventPublisher,
                               @Nonnull StreamObserver<RpcHostingWriteStreamFragment> observer,
                               @Nonnull UpdateRefsContextHolder updateRefsContextHolder) {
            super(POST_RECEIVE, timeout);
            this.eventPublisher = requireNonNull(eventPublisher, "eventPublisher");
            this.maxFragmentSize = maxFragmentSize;
            this.observer = requireNonNull(observer, "observer");
            this.updateRefsContextHolder = updateRefsContextHolder;
        }

        @Override
        public boolean handle(@Nonnull Repository repository, @Nonnull HookRequest request,
                              @Nonnull GitHookOutputWriter writer) {
            //Signal that one or more refs have been updated
            eventPublisher.publish(new RefsUpdatedEvent(repository));

            Iterator<List<RpcRefUpdate>> fragmentIt = parseRefUpdates(request.getStdin(), maxFragmentSize).iterator();
            while (fragmentIt.hasNext()) {
                List<RpcRefUpdate> refUpdates = fragmentIt.next();
                observer.onNext(RpcHostingWriteStreamFragment.newBuilder()
                        .setPostReceive(RpcHostingPostReceiveRequest.newBuilder()
                                .addAllRefUpdates(refUpdates)
                                .setMoreRefUpdates(fragmentIt.hasNext()))
                        .build());

                UpdateRefsContext updateRefsContext = updateRefsContextHolder.get();
                if (updateRefsContext != null) {
                    updateRefsContext.clearCommitted(refUpdates.stream()
                            .map(update -> update.getRefId().toStringUtf8())
                            .collect(toList()));
                }
            }
            return super.handle(repository, request, writer);
        }

        @Override
        protected boolean isAccepted(@Nonnull RpcHostingPostReceiveResponse response) {
            return true;
        }

        @Nonnull
        @Override
        protected List<RpcHookChunk> readOutput(@Nonnull RpcHostingPostReceiveResponse response) {
            return response.getOutputList();
        }
    }

    private static class PreReceiveHookHandler extends CompletableGitHookHandler<RpcHostingPreReceiveResponse> {

        private final long maxFragmentSize;
        private final StreamObserver<RpcHostingWriteStreamFragment> observer;
        private final GitQuarantineRegistry quarantineRegistry;
        private final UpdateRefsContextHolder updateRefsContextHolder;

        PreReceiveHookHandler(GitQuarantineRegistry quarantineRegistry, UpdateRefsContextHolder updateRefsContextHolder,
                              long maxFragmentSize, Duration timeout,
                              StreamObserver<RpcHostingWriteStreamFragment> observer) {
            super(PRE_RECEIVE, timeout);
            this.maxFragmentSize = maxFragmentSize;
            this.observer = observer;
            this.quarantineRegistry = quarantineRegistry;
            this.updateRefsContextHolder = updateRefsContextHolder;
        }

        @Override
        public boolean handle(@Nonnull Repository repository, @Nonnull HookRequest request,
                              @Nonnull GitHookOutputWriter writer) {
            String quarantinedRepositoryId = registerQuarantineEnvironment(repository, request.getEnvironment());
            List<List<RpcRefUpdate>> refUpdateFragments = parseRefUpdates(request.getStdin(), maxFragmentSize);
            Iterator<List<RpcRefUpdate>> fragmentIt = refUpdateFragments.iterator();
            while (fragmentIt.hasNext()) {
                List<RpcRefUpdate> refUpdates = fragmentIt.next();
                observer.onNext(RpcHostingWriteStreamFragment.newBuilder()
                        .setPreReceive(RpcHostingPreReceiveRequest.newBuilder()
                                .addAllRefUpdates(refUpdates)
                                .setMoreRefUpdates(fragmentIt.hasNext())
                                .setRepository(quarantinedRepositoryId))
                        .build());
            }
            if (super.handle(repository, request, writer)) {
                updateRefsContextHolder.set(repository, refUpdateFragments.stream()
                        .flatMap(Collection::stream)
                        .map(RefUpdate::of)
                        .collect(toList()));
                return true;
            }
            return false;
        }

        @Override
        protected boolean isAccepted(@Nonnull RpcHostingPreReceiveResponse response) {
            return !response.getRejected();
        }

        @Nonnull
        @Override
        protected List<RpcHookChunk> readOutput(@Nonnull RpcHostingPreReceiveResponse response) {
            return response.getOutputList();
        }

        @Nonnull
        private String registerQuarantineEnvironment(Repository repository, Map<String, String> environment) {
            return GitQuarantineEnvironment.fromEnvironment(repository, environment)
                    .map(quarantineRegistry::register)
                    .orElse(repository.getId());
        }
    }

    private static class RefChangeFragmentChunker {

        private final ImmutableList.Builder<List<RpcRefUpdate>> builder = ImmutableList.builder();
        private final long maxFragmentSize;
        private ImmutableList.Builder<RpcRefUpdate> fragmentBuilder = ImmutableList.builder();
        private long fragmentSize;

        RefChangeFragmentChunker(long maxFragmentSize) {
            this.maxFragmentSize = maxFragmentSize;
        }

        void addRefUpdate(String refId, String oldHash, String newHash) {
            RpcRefUpdate refUpdate = RpcRefUpdate.newBuilder()
                    .setRefId(ByteString.copyFromUtf8(refId))
                    .setNewHash(newHash)
                    .setOldHash(oldHash)
                    .build();
            long size = refUpdate.getSerializedSize();
            if (fragmentSize + size > maxFragmentSize) {
                builder.add(fragmentBuilder.build());
                fragmentBuilder = ImmutableList.builder();
                fragmentSize = 0L;
            }
            fragmentBuilder.add(refUpdate);
            fragmentSize += size;
        }

        List<List<RpcRefUpdate>> getResult() {
            builder.add(fragmentBuilder.build());
            fragmentBuilder = ImmutableList.builder();
            return builder.build();
        }
    }

    private static class RefChangeLineProcessor extends AbstractRefUpdateParser<List<List<RpcRefUpdate>>> {

        private final RefChangeFragmentChunker refChangeFragmentChunker;

        RefChangeLineProcessor(RefChangeFragmentChunker refChangeFragmentChunker) {
            this.refChangeFragmentChunker = refChangeFragmentChunker;
        }

        @Nonnull
        @Override
        public List<List<RpcRefUpdate>> getResult() {
            return refChangeFragmentChunker.getResult();
        }

        @Override
        protected boolean processRefUpdate(String refId, String oldHash, String newHash) {
            refChangeFragmentChunker.addRefUpdate(refId, oldHash, newHash);
            return true;
        }
    }
}
