package com.atlassian.bitbucket.mesh.grpc;

import com.atlassian.bitbucket.mesh.process.OsProcess;
import com.atlassian.bitbucket.mesh.process.ProcessCompletionHandler;
import com.atlassian.bitbucket.mesh.process.ProcessStdioHandler;
import com.atlassian.bitbucket.mesh.process.ProcessSummary;
import com.atlassian.bitbucket.mesh.rpc.v1.RpcStreamChunk;
import com.atlassian.bitbucket.mesh.rpc.v1.RpcStreamFragment;
import com.zaxxer.nuprocess.NuProcess;
import io.grpc.Status;
import io.grpc.stub.ClientCallStreamObserver;
import io.grpc.stub.ClientResponseObserver;
import org.apache.commons.lang3.StringUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.annotation.Nonnull;
import javax.annotation.Nullable;
import java.nio.ByteBuffer;
import java.nio.charset.StandardCharsets;
import java.util.concurrent.atomic.AtomicBoolean;

import static java.util.Objects.requireNonNull;

/**
 * A base class for constructing {@link ClientResponseObserver} implementations which receive {@link RpcStreamFragment}s
 * and forward them to an {@link ProcessStdioHandler}
 * <p>
 * This class contains a lot of similar logic to {@code NioNuProcessHandler}, for handling {@code stderr} and
 * invoking a {@link ProcessCompletionHandler} to evaluate the request's outcome after receiving the exit code.
 * <p>
 * This base class provides {@code stdout} and {@code stderr} handling suitable for use in server-side streaming or
 * bidirectional streaming requests. For bidirectional streaming, subclasses must provide their own {@code stdin}
 * handling.
 *
 * @param <ReqT> the request type for sending data to the remote process
 * @param <T>    the type of result returned by the handler, if the request completes successfully
 */
public abstract class AbstractFragmentResponseObserver<ReqT, T>
        extends AbstractFutureResponseObserver<ReqT, RpcStreamFragment, T> {

    protected final Logger log;
    protected final ProcessStdioHandler<T> stdioHandler;

    private final ProcessCompletionHandler completionHandler;
    private final StderrCallback stderrCallback;
    private final StdoutCallback stdoutCallback;

    private int exitCode;
    private RemoteProcess<ReqT> process;
    private HandlerState state;
    private StringBuilder stderrBuffer;
    private Throwable thrown;

    /**
     * @param stdioHandler the I/O <i>and completion</i> handler
     * @param <H> the handler's type
     */
    protected <H extends ProcessStdioHandler<T> & ProcessCompletionHandler> AbstractFragmentResponseObserver(
            @Nonnull H stdioHandler) {
        this(stdioHandler, null);
    }

    /**
     * @param stdioHandler      the I/O handler, which may also be a completion handler
     * @param completionHandler the completion handler, if the I/O handler is not one
     */
    protected AbstractFragmentResponseObserver(@Nonnull ProcessStdioHandler<T> stdioHandler,
                                               @Nullable ProcessCompletionHandler completionHandler) {
        this.completionHandler = chooseCompletionHandler(stdioHandler, completionHandler);
        this.stdioHandler = requireNonNull(stdioHandler, "stdioHandler");

        exitCode = -1002;
        log = LoggerFactory.getLogger(getClass());
        state = HandlerState.CREATED;
        stderrCallback = new StderrCallback();
        stdoutCallback = new StdoutCallback();
    }

    @Override
    public T asResult() {
        //Block until completion. The return value will always be null, though, so we ignore it
        super.asResult();

        //Call the completion handler to allow it to process the exit code, thrown, etc.
        callCompletionHandler();

        //If we make it here, the completion handler didn't throw, so we're good to return the result.
        return stdioHandler.getResult();
    }

    @Override
    public void beforeStart(ClientCallStreamObserver<ReqT> requestStream) {
        super.beforeStart(requestStream);

        //Wrap the stream in an NioProcess to provide to handlers, to allow them to cancel
        process = createProcess(requestStream);

        //Note: beforeStart is called before gRPC starts the remote call, so we don't call onStart here
        //      because we don't know whether the process has actually been started on the remote side
        invokeCallback(() -> stdioHandler.onPreStart(process));
    }

    @Override
    public void onCompleted() {
        //No-op; finish is called from onNext when it receives the exit code or when onError is called,
        //and it always calls either onCompleted or onError on the superclass
    }

    @Override
    public void onError(Throwable t) {
        //Check for some common statuses
        Status.Code code = Status.fromThrowable(t).getCode();
        if (code == Status.Code.CANCELLED) {
            if (isCanceled()) {
                //If we're already canceled, we ignore the cancellation exception (since it's ours)
                //TODO: Is remote cancellation going to be a thing, for streaming operations?
                t = null;
                if (exitCode == -1002) {
                    stdioHandler.onExit(-1);
                }
            }
        }

        if (t != null) {
            setOrUpdateThrown(t);
        }

        //onNext with an RpcStreamResult may never be called, depending on what went wrong, so we need
        //to call finish here to ensure we wrap up the handler
        finish(-1001);
    }

    @Override
    public void onNext(RpcStreamFragment fragment) {
        switch (fragment.getFragmentOneofCase()) {
            case RESULT:
                callOnExit(fragment.getResult().getExitCode());
                break;
            case START:
                callOnStart(fragment.getStart().getCommandLine().toStringUtf8());
                break;
            case STDERR:
                stderrCallback.onChunk(fragment.getStderr());
                break;
            case STDOUT:
                stdoutCallback.onChunk(fragment.getStdout());
                break;
        }
    }

    /**
     * Creates a {@link RemoteProcess} wrapping the provided {@link ClientCallStreamObserver}, wrapping gRPC behaviors
     * in an {@link OsProcess} implementation that I/O handlers can interact with. This allows reusing the existing
     * handlers, but backed by gRPC rather than a local process.
     *
     * @param requestStream the stream
     * @return a process backed by the provided stream
     */
    protected abstract RemoteProcess<ReqT> createProcess(ClientCallStreamObserver<ReqT> requestStream);

    @Override
    protected boolean isCanceled() {
        return super.isCanceled() || (process != null && process.isCanceled());
    }

    /**
     * Returns the provided {@link ProcessStdioHandler} if it also implements {@link ProcessCompletionHandler}, or
     * returns the provided {@link ProcessCompletionHandler}.
     * <p>
     * This parallels {@code DefaultGitProcessBuilder.chooseCompletionHandler}'s logic.
     *
     * @param stdioHandler      the I/O handler, which may also be a completion handler
     * @param completionHandler the completion handler, which may be {@code null} if the I/O handler is also a
     *                          completion handler
     * @return the completion handler to use
     * @throws NullPointerException if the I/O handler isn't a completion handler and no separate completion
     *                              handler was supplied
     */
    private static ProcessCompletionHandler chooseCompletionHandler(ProcessStdioHandler<?> stdioHandler,
                                                                    ProcessCompletionHandler completionHandler) {
        //If the I/O handler is also a completion handler, that supersedes any separate
        //completion handler. Otherwise, a separate completion handler is required. This
        //parallels how DefaultGitProcessBuilder.chooseCompletionHandler works
        if (stdioHandler instanceof ProcessCompletionHandler) {
            return (ProcessCompletionHandler) stdioHandler;
        }

        return requireNonNull(completionHandler, "completionHandler");
    }

    private void callCompletionHandler() {
        String commandLine = process == null ? "<Unknown>" : process.toString();
        String stderr = stderrBuffer == null ? null : StringUtils.chomp(stderrBuffer.toString());

        completionHandler.onComplete(new ProcessSummary.Builder(commandLine)
                .canceled(isCanceled())
                .exitCode(exitCode)
                .stderr(stderr)
                .thrown(thrown)
                .timedOut(thrown != null && Status.fromThrowable(thrown).getCode() == Status.Code.DEADLINE_EXCEEDED)
                .build());
    }

    private void callOnExit(int exitCode) {
        process.onExit();

        //Intentionally not using invokeCallback; we don't want the maybeAbort functionality here
        try {
            stdioHandler.onExit(exitCode);
        } catch (Throwable t) { //Yes, Throwable. We need to catch Errors too here
            //Exceptions from onExit take precedence over other exceptions thrown by callbacks
            //while the process was running, so this doesn't call setOrUpdateThrown; we always
            //want to assign thrown here
            if (thrown != null) {
                t.addSuppressed(thrown);
            }
            thrown = t;
        }

        finish(exitCode);
    }

    private void callOnStart(String commandLine) {
        process.onStart(commandLine);
        state = HandlerState.STARTED;

        invokeCallback(() -> stdioHandler.onStart(process));
    }

    private void ensureStarted() {
        if (state == HandlerState.CREATED) {
            IllegalStateException e = new IllegalStateException("The remote process could not be started");
            if (thrown == null) { //onPreStart could have thrown and set an exception already
                thrown = e;
            } else {
                thrown.addSuppressed(e);
            }
        }
    }

    private void finish(int code) {
        if (state == HandlerState.FINISHED) {
            return;
        }

        exitCode = code;
        try {
            //This must be done before the state is updated; otherwise it can't detect whether onStart was
            //called or not because the state will be FINISHED
            ensureStarted();

            //Always complete the future normally. Exceptions, if any, are handled in on the request thread
            //in asResult() so their stack traces will be more useful
            super.onCompleted();
        } finally {
            state = HandlerState.FINISHED;
        }
    }

    private void invokeCallback(Runnable callback) {
        try {
            callback.run();
        } catch (Throwable t) { //Yes, Throwable. We need to catch Errors too in this very special case
            setOrUpdateThrown(t);

            //If one of our handlers fails and the process is still running, terminate it
            log.debug("[{}]: Destroying after handler failure", process, t);
            process.abort();
        }
    }

    private boolean maybeAbort() {
        if (thrown == null) {
            return false;
        }

        log.debug("[{}]: Dropping callback invocation", process, thrown);
        process.abort();

        return true;
    }

    private void setOrUpdateThrown(Throwable t) {
        if (thrown == null) {
            thrown = t;
        } else {
            thrown.addSuppressed(t);
        }
    }

    /**
     * Describes an {@link OsProcess} implementation which is backed by a remote process.
     * <p>
     * This base class provides cancellation functionality, as well as a basic {@code toString()} implementation
     * which returns the process's command line, if the process {@link #isStarted is started}.
     *
     * @param <ReqT> the request type for sending data to the remote process
     */
    protected abstract static class RemoteProcess<ReqT> implements OsProcess {

        protected final ClientCallStreamObserver<ReqT> requestStream;

        private final AtomicBoolean canceled;

        private volatile String commandLine;
        private volatile boolean running;

        RemoteProcess(ClientCallStreamObserver<ReqT> requestStream) {
            this.requestStream = requestStream;

            canceled = new AtomicBoolean();
        }

        @Override
        public void cancel() {
            if (canceled.compareAndSet(false, true)) {
                requestStream.cancel("The request was canceled by a handler", null);
            }
        }

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

        @Override
        public boolean isRunning() {
            return running;
        }

        @Override
        public String toString() {
            return commandLine == null ? "<Not started>" : commandLine;
        }

        protected boolean isStarted() {
            return commandLine != null;
        }

        protected void onExit() {
            running = false;
        }

        protected void onStart(@Nonnull String commandLine) {
            this.commandLine = commandLine;

            running = true;
        }

        void abort() {
            if (canceled.compareAndSet(false, true)) {
                requestStream.cancel("The request was aborted after a handler failure", null);
            }
        }
    }

    private abstract class AbstractChunkCallback {

        private ByteBuffer proxy;

        public void onChunk(RpcStreamChunk chunk) {
            ByteBuffer buffer = chunk.getData().asReadOnlyByteBuffer();
            boolean closed = chunk.getClosed();

            if (proxy == null) {
                //If we don't have a proxy either this is the first callback or the handler consumed the
                //full buffer during the previous callback
                handleOutput(buffer, closed);

                if (buffer.hasRemaining() && shouldProxy(buffer, closed)) { //stderr may not be proxyable
                    //Note: We use NuProcess's buffer capacity here because that's the buffer size Mesh
                    //will be using on the remote side.
                    proxy = ByteBuffer.allocate(NuProcess.BUFFER_CAPACITY);
                    proxy.put(buffer);
                }
            } else {
                //If we have a proxy, it means the handler didn't fully consume the buffer on a previous
                //call. That means we need to stream subsequent buffers "through" the proxy, until we get
                //to a state where the proxy buffer is empty
                pumpViaProxy(buffer, closed);
            }
        }

        protected abstract void accept(ByteBuffer buffer, boolean closed);

        protected boolean shouldProxy(ByteBuffer buffer, boolean closed) {
            return !(closed || isCanceled());
        }

        private void handleOutput(ByteBuffer buffer, boolean closed) {
            invokeCallback(() -> {
                if (maybeAbort()) {
                    //If processing is "done" (generally either meaning it has failed, timed out or been
                    //canceled, don't pump any more output to the callback
                    buffer.position(buffer.position() + buffer.remaining());
                } else {
                    accept(buffer, closed);
                }
            });
        }

        private void pumpViaProxy(ByteBuffer buffer, boolean closed) {
            int remaining = buffer.remaining();
            //Use a do/while loop so that, if the _new_ buffer is empty but the _proxy_ buffer isn't, we forward those
            //proxy contents on to the handler.
            do {
                int available = proxy.remaining();
                if (remaining > available) {
                    ByteBuffer slice = buffer.slice();
                    slice.limit(available);

                    proxy.put(slice);
                    buffer.position(buffer.position() + available); //Advance the buffer
                    remaining -= available;
                } else {
                    proxy.put(buffer);
                    remaining = 0;
                }
                proxy.flip(); //Flip for reading

                //Only pass true for closed if we've fully consumed the remote buffer _and_ the chunk indicates
                //the remote stream is closed
                handleOutput(proxy, remaining == 0 && closed);

                if (isCanceled() || (remaining == 0 && !proxy.hasRemaining())) {
                    //If the request is canceled, or we get to a point where both the remote and proxy buffers
                    //are empty, it means we can try to consume the remote buffer directly on the next chunk.
                    //For performance reasons, if we can get back to that state we _want_ to (to reduce copying)
                    proxy = null;
                    break;
                } else {
                    //Otherwise, if the remote buffer isn't empty or the proxy buffer wasn't drained, compact it
                    //to prepare for the next loop (or callback)
                    proxy.compact();
                }
            } while (remaining > 0); //The remote buffer may spread into multiple callbacks
        }
    }

    private enum HandlerState {

        CREATED,
        FINISHED,
        STARTED
    }

    private class StderrCallback extends AbstractChunkCallback {

        @Override
        protected void accept(ByteBuffer buffer, boolean closed) {
            stdioHandler.onStderr(buffer, closed);
        }

        @Override
        protected boolean shouldProxy(ByteBuffer buffer, boolean closed) {
            if (buffer.position() == 0) {
                //If the buffer isn't empty and the delegate handler didn't consume any bytes, assume they
                //don't want to and buffer stderr to pass to the exit handler
                if (stderrBuffer == null) {
                    stderrBuffer = new StringBuilder();
                }
                //TODO: If stderr is split across multiple ByteBuffers, does this work correctly?
                stderrBuffer.append(StandardCharsets.UTF_8.decode(buffer));

                return false; //Don't create a proxy; we consumed the buffer
            }

            return super.shouldProxy(buffer, closed);
        }
    }

    private class StdoutCallback extends AbstractChunkCallback {

        @Override
        protected void accept(ByteBuffer buffer, boolean closed) {
            stdioHandler.onStdout(buffer, closed);
        }
    }
}
