package com.atlassian.bitbucket.mesh.git;

import com.atlassian.bitbucket.mesh.git.hook.script.CallHookScriptParameters;
import com.atlassian.bitbucket.mesh.git.hook.script.HookScriptCallback;
import com.atlassian.bitbucket.mesh.git.hook.script.HookScriptManager;
import com.atlassian.bitbucket.mesh.git.ref.RefUpdate;
import com.atlassian.bitbucket.mesh.process.ProcessSummary;
import com.atlassian.bitbucket.mesh.repository.Repository;
import com.atlassian.bitbucket.mesh.repository.RepositoryManager;
import com.atlassian.bitbucket.mesh.rpc.v1.*;
import com.google.protobuf.ByteString;
import io.grpc.Status;
import io.grpc.stub.StreamObserver;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.annotation.Nonnull;
import java.nio.ByteBuffer;
import java.time.Duration;
import java.util.HashSet;
import java.util.List;
import java.util.Set;
import java.util.concurrent.TimeUnit;
import java.util.stream.Collectors;

import static com.atlassian.bitbucket.mesh.rpc.util.ByteStringUtils.toStringsLazily;
import static com.atlassian.bitbucket.mesh.rpc.v1.RpcHookChunk.Channel.CHANNEL_STDERR;
import static com.atlassian.bitbucket.mesh.rpc.v1.RpcHookChunk.Channel.CHANNEL_STDOUT;
import static com.atlassian.bitbucket.mesh.rpc.v1.RpcHookResult.HOOK_RESULT_ACCEPTED;
import static com.atlassian.bitbucket.mesh.rpc.v1.RpcHookResult.HOOK_RESULT_REJECTED;

public class CallHookScriptsRequestObserver implements StreamObserver<RpcCallHookScriptsFragment> {

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

    private final HookScriptManager hookScriptManager;
    private final RepositoryManager repositoryManager;
    private final StreamObserver<RpcCallHookScriptsResponseFragment> responseObserver;

    private volatile boolean canceled;
    private volatile boolean completed;
    private Set<String> missingScripts;
    private RpcCallHookScriptsRequest request;

    public CallHookScriptsRequestObserver(HookScriptManager hookScriptManager, RepositoryManager repositoryManager,
                                          StreamObserver<RpcCallHookScriptsResponseFragment> responseObserver) {
        this.hookScriptManager = hookScriptManager;
        this.repositoryManager = repositoryManager;
        this.responseObserver = responseObserver;
    }

    @Override
    public void onError(Throwable t) {
        if (Status.fromThrowable(t).getCode() == Status.Code.CANCELLED) {
            //Log cancellation more quietly; it's not an error for the client to cancel
            log.debug("The client has canceled the request");
        } else if (request == null) {
            log.warn("The call failed before a request was sent", t);
        } else if (completed) {
            log.warn("The call failed, but the requested process has already completed", t);
        } else {
            log.error("Not calling hook scripts; the client has failed", t);
        }
        canceled = true;
        completed = true;
    }

    @Override
    public void onNext(RpcCallHookScriptsFragment fragment) {
        switch (fragment.getFragmentOneofCase()) {
            case REQUEST:
                request = fragment.getRequest();
                missingScripts = getMissingScripts(request.getHookScriptList());
                if (!missingScripts.isEmpty()) {
                    responseObserver.onNext(RpcCallHookScriptsResponseFragment.newBuilder()
                            .setMissing(RpcHookScriptsMissing.newBuilder()
                                    .addAllHookScript(missingScripts)
                                    .build())
                            .build());
                }
                break;
            case HOOK_SCRIPT:
                RpcHookScript hookScript = fragment.getHookScript();
                String hookScriptId = hookScript.getId();
                hookScriptManager.put(hookScriptId, outputStream -> hookScript.getContent().writeTo(outputStream));
                missingScripts.remove(hookScriptId);
                break;
        }

        if (missingScripts.isEmpty()) {
            //All requested scripts are available locally, time to execute them!
            callScripts();
        }
    }

    @Override
    public void onCompleted() {
        //Client has sent us all they're going to send, even if we haven't received all missing scripts, invoke
        //what we can.
        if (!completed) {
            if (!missingScripts.isEmpty()) {
                log.warn("Client sent onCompleted, but some hook scripts have not been received: {}", missingScripts);
            }

            //Any missing hook scripts will be detected by HookScriptManager.call and will be marked as rejected
            callScripts();
        }
    }

    private void callScripts() {
        if (completed) {
            return;
        }

        Repository repository = repositoryManager.getById(request.getRepository());
        RpcHookScriptOptions options = request.getHookScriptOptions();
        CallHookScriptParameters parameters = new CallHookScriptParameters.Builder(repository)
                .abortOnFirstVeto(request.getAbortOnFirstVeto())
                .environment(toStringsLazily(options.getEnvironmentMap()))
                .executionTimeout(Duration.ofSeconds(options.getTimeouts().getExecution()))
                .idleTimeout(Duration.ofSeconds(options.getTimeouts().getIdle()))
                .refUpdates(request.getRefUpdatesList().stream()
                        .map(RefUpdate::of)
                        .collect(Collectors.toList()))
                .build();

        RpcHookScriptCallback callback = new RpcHookScriptCallback(options.getMaxOutputLength());
        hookScriptManager.call(request.getHookScriptList(), parameters, callback);
        RpcCallHookScriptsResponse response = callback.getResponse();
        responseObserver.onNext(RpcCallHookScriptsResponseFragment.newBuilder()
                .setResponse(response)
                .build());
        responseObserver.onCompleted();
        completed = true;
        log.debug("[{}] Completed calling {} hook scripts", repository, response.getResponsesCount());
    }

    private Set<String> getMissingScripts(List<String> hookScriptIds) {
        Set<String> missing = new HashSet<>();
        for (String id : hookScriptIds) {
            if (!hookScriptManager.exists(id)) {
                missing.add(id);
            }
        }
        return missing;
    }

    private class RpcHookScriptCallback implements HookScriptCallback {

        private final RpcCallHookScriptsResponse.Builder builder;
        private final int maxBytes;
        private final Object lock;

        private int bytesRemaining;
        private ByteString currentBytes;
        private RpcHookChunk.Channel currentChannel;
        private String currentScript;
        private long lastEmittedTimestamp;
        private RpcHookScriptResponse.Builder scriptBuilder;
        private long startNanos;

        RpcHookScriptCallback(int maxBytes) {
            this.maxBytes = maxBytes;

            bytesRemaining = maxBytes;
            builder = RpcCallHookScriptsResponse.newBuilder();
            lock = new Object();
        }

        @Override
        public boolean isCanceled() {
            return canceled;
        }

        @Override
        public void onComplete(@Nonnull ProcessSummary processSummary) {
            emitCurrentChunk();

            scriptBuilder.setTimedOut(processSummary.isTimedOut());
            scriptBuilder.setResult(processSummary.getExitCode() == 0 ? HOOK_RESULT_ACCEPTED : HOOK_RESULT_REJECTED);
            scriptBuilder.setDurationMs(TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - startNanos));
            builder.addResponses(scriptBuilder.build());
            scriptBuilder = null;
        }

        @Override
        public void onStartScript(@Nonnull String hookScriptId) {
            scriptBuilder = RpcHookScriptResponse.newBuilder().setHookScript(hookScriptId);
            bytesRemaining = maxBytes;
            startNanos = System.nanoTime();
            currentScript = hookScriptId;
        }

        @Override
        public void onStderr(@Nonnull ByteBuffer buffer, boolean closed) {
            onBytes(buffer, CHANNEL_STDERR);
        }

        @Override
        public void onStdout(@Nonnull ByteBuffer buffer, boolean closed) {
            onBytes(buffer, CHANNEL_STDOUT);
        }

        RpcCallHookScriptsResponse getResponse() {
            emitCurrentChunk();
            return builder.build();
        }

        private void emitCurrentChunk() {
            synchronized (lock) {
                if (currentBytes != null && currentChannel != null) {
                    responseObserver.onNext(RpcCallHookScriptsResponseFragment.newBuilder()
                            .setChunk(RpcHookScriptChunk.newBuilder()
                                    .setHookScript(currentScript)
                                    .setChunk(RpcHookChunk.newBuilder()
                                            .setData(currentBytes)
                                            .setChannel(currentChannel)))
                            .build());
                    lastEmittedTimestamp = System.currentTimeMillis();
                }
                currentBytes = null;
                currentChannel = null;
            }
        }

        private void onBytes(ByteBuffer buffer, RpcHookChunk.Channel channel) {
            int size = Math.min(buffer.remaining(), bytesRemaining);
            if (size > 0) {
                ByteString bytes = ByteString.copyFrom(buffer, size);
                synchronized (lock) {
                    if (currentChannel == channel) {
                        currentBytes = currentBytes.concat(bytes);
                        if (System.currentTimeMillis() - lastEmittedTimestamp > 3000) {
                            emitCurrentChunk();
                        }
                    } else {
                        emitCurrentChunk();

                        currentChannel = channel;
                        currentBytes = bytes;
                    }
                }
            }
            buffer.position(buffer.position() + buffer.remaining());
        }
    }
}
