package com.atlassian.bitbucket.mesh.grpc;

import com.atlassian.bitbucket.mesh.git.exception.AbortedException;
import com.atlassian.bitbucket.mesh.process.ProcessCompletionHandler;
import com.atlassian.bitbucket.mesh.process.ProcessStdioHandler;
import com.atlassian.bitbucket.mesh.rpc.v1.RpcStreamChunk;
import com.google.protobuf.ByteString;
import com.zaxxer.nuprocess.NuProcess;
import io.grpc.Status;
import io.grpc.stub.ClientCallStreamObserver;
import io.grpc.stub.StreamObserver;

import javax.annotation.Nonnull;
import javax.annotation.Nullable;
import java.nio.ByteBuffer;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.atomic.AtomicBoolean;

import static java.util.Objects.requireNonNull;

/**
 * A {@link StreamObserver} implementation which supports <i>bidirectional</i> streaming. An {@link ProcessStdioHandler}
 * is required to allow providing {@code stdin}.
 * <p>
 * Because gRPC/Protobuf don't allow inheritance or interfaces or other mechanisms to allow for polymorphism in
 * messages, a secondary {@link RequestFragmentCreator} is used to provide an abstraction over how fragments are
 * created.
 *
 * @param <ReqT> the fragment request type
 * @param <T>    the type of result returned by the handler, if the request completes successfully
 */
public class BidirectionalFragmentResponseObserver<ReqT, T> extends AbstractFragmentResponseObserver<ReqT, T> {

    private final RequestFragmentCreator<ReqT> fragmentCreator;

    private BidirectionalRemoteProcess process;

    /**
     * Creates a new {@code BidirectionalFragmentResponseObserver} which uses a {@link ProcessStdioHandler} that
     * also implements {@link ProcessCompletionHandler}.
     *
     * @param stdioHandler    the I/O <i>and completion</i> handler
     * @param fragmentCreator the fragment creator
     * @param <H> the handler's type
     */
    public <H extends ProcessStdioHandler<T> & ProcessCompletionHandler> BidirectionalFragmentResponseObserver(
            @Nonnull H stdioHandler, @Nonnull RequestFragmentCreator<ReqT> fragmentCreator) {
        this(stdioHandler, null, fragmentCreator);
    }

    /**
     * Creates a new {@code BidirectionalFragmentResponseObserver} which uses the provided {@link ProcessStdioHandler}
     * and {@link ProcessCompletionHandler}. If the provided I/O handler implements {@link ProcessCompletionHandler}
     * the separate handler may be {@code null}; otherwise, an explicit completion handler must be provided.
     *
     * @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
     * @param fragmentCreator   the fragment creator
     */
    public BidirectionalFragmentResponseObserver(@Nonnull ProcessStdioHandler<T> stdioHandler,
                                                 @Nullable ProcessCompletionHandler completionHandler,
                                                 @Nonnull RequestFragmentCreator<ReqT> fragmentCreator) {
        super(stdioHandler, completionHandler);

        this.fragmentCreator = requireNonNull(fragmentCreator, "fragmentCreator");
    }

    @Override
    protected RemoteProcess<ReqT> createProcess(ClientCallStreamObserver<ReqT> requestStream) {
        return process = new BidirectionalRemoteProcess(requestStream);
    }

    @Override
    public T asResult() {
        //First, add a listener so we can notify the process when the request completes. This will
        //allow the run() loop to terminate
        addListener((result, thrown) -> process.onComplete());

        //Run the process, which will pump stdin as necessary _on this thread_. This avoids creating
        //a separate thread to handle that
        process.run();

        //When run() returns, the request should be complete and the result should be available
        return super.asResult();
    }

    /**
     * A bidirectional {@link RemoteProcess} implementation which allows writing {@code stdin} fragments.
     */
    private class BidirectionalRemoteProcess extends RemoteProcess<ReqT> implements Runnable {

        private final AtomicBoolean closeRequested;
        private final AtomicBoolean closeSent;
        private final ByteBuffer stdin;
        private final BlockingQueue<Runnable> tasks;
        private final AtomicBoolean writeRequested;

        BidirectionalRemoteProcess(ClientCallStreamObserver<ReqT> requestStream) {
            super(requestStream);

            closeRequested = new AtomicBoolean();
            closeSent = new AtomicBoolean();
            stdin = ByteBuffer.allocate(NuProcess.BUFFER_CAPACITY);
            tasks = new LinkedBlockingQueue<>();
            writeRequested = new AtomicBoolean();

            requestStream.setOnReadyHandler(this::onReady);
        }

        @Override
        public void closeStdin(boolean force) {
            if (force) {
                //Forced closes are always sent to the server immediately, unless we've already
                //sent a close through previously
                closeRequested.set(true);
                if (closeSent.compareAndSet(false, true)) {
                    //This calls onNext immediately, even if it may end up buffered
                    callOnNext(ByteString.EMPTY, true);
                }
            } else if (closeRequested.compareAndSet(false, true)) {
                //Otherwise, the first time closeStdin is called we queue a task to send the
                //close when the observer is ready
                tasks.offer(() -> {
                    if (closeSent.compareAndSet(false, true)) {
                        //This uses sendStdin/requireReady to prevent buffering
                        sendStdin(ByteString.EMPTY, true);
                    }
                });
            }
        }

        @Override
        public boolean hasPendingWrites() {
            //Nothing actually uses this at the moment, so a "correct" implementation doesn't
            //necessarily matter, but if anything tries to _start_ using this it would be nice
            //if it worked. The current code is an approximation; if the requestStream isn't
            //ready then we know there are already one or more fragments pending write
            return !requestStream.isReady();
        }

        /**
         * Dispatches the initial request fragment and then processes {@code stdin} on the client to dispatch
         * it to the server.
         * <p>
         * Methods like {@link #closeStdin} and {@link #wantWrite} can be called from different threads, and
         * {@link #onReady} is always invoked from a gRPC executor thread. The delegate {@link #stdioHandler}
         * is not thread-safe, however, so we need to ensure it's only used by a single thread. Additionally,
         * we <i>must not</i> block the gRPC thread that calls {@link #onReady}, as doing so can cause request
         * processing to deadlock.
         * <p>
         * One option would be to spin up a separate executor thread, but serializing access with that approach
         * is problematic--if we submit multiple tasks to the executor they could run concurrently, and if we
         * "simply" added synchronization to them then we would be wasting executor threads. Rather than doing
         * that, we leverage a thread we already have: The thread that initiated the request. That thread would
         * otherwise be blocked waiting for the request to complete, so it should be safe for us to use it. As
         * a bonus, using the initiating thread also serves to ensure only a single thread interacts with the
         * {@link #stdioHandler} without requiring any extra synchronization or blocking other threads.
         */
        @Override
        public void run() {
            requestStream.onNext(fragmentCreator.getInitial());

            try {
                //The blocking take() here should be safe; even if the request is canceled or
                //times out onComplete() will be called when the request's future completes
                //and release this thread (otherwise waiting for asResult() would hang too)
                Runnable task;
                while ((task = tasks.take()) != this) {
                    try {
                        task.run();
                    } catch (Exception e) {
                        //Forward any exception we get to the server to terminate the request
                        //TODO: Is this correct?
                        requestStream.onError(e);
                        break;
                    }
                }
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt();

                //TODO: I'm not sure what to do here. This may not be correct
                requestStream.onError(Status.ABORTED
                        .withDescription("Interrupted while waiting for the request to complete")
                        .asRuntimeException());
            }
        }

        /**
         * Sets {@link #writeRequested} and then queues a task for the handler thread (where the {@link #run}
         * loop is executing) to {@link #runStdinLoop read from stdin}. The actual read is not performed on
         * the calling thread to ensure the delegate {@link #stdioHandler} is not accessed concurrently. This
         * also mirrors how NuProcess would work, since {@code wantWrite()} would trigger an {@code onStdinReady}
         * callback on the thread pumping the process rather than the requesting thread.
         */
        @Override
        public void wantWrite() {
            if (!isStarted()) {
                throw new IllegalStateException("Wait until onStart to call wantWrite");
            }
            if (closeRequested.get()) {
                throw new IllegalStateException("wantWrite cannot be called after stdin is closed");
            }

            writeRequested.set(true);
            tasks.offer(this::runStdinLoop);
        }

        /**
         * Writes the provided {@code buffer} directly to {@link StreamObserver#onNext} <i>even if doing so will
         * cause gRPC to buffer the message</i>.
         * <p>
         * NuProcess provides no flow control for {@code writeStdin}, and that must be emulated here. Blocking to
         * wait for {@link ClientCallStreamObserver#isReady} may deadlock request processing, if {@code writeStdin}
         * is called on a gRPC thread, because it will prevent {@link #onReady} from being called.
         *
         * @param buffer a buffer containing one or more bytes to write to {@code stdin}
         */
        @Override
        public void writeStdin(@Nonnull ByteBuffer buffer) {
            if (!isStarted()) {
                throw new IllegalStateException("Wait until onStart to call writeStdin");
            }
            if (closeRequested.get()) {
                throw new IllegalStateException("writeStdin cannot be called after stdin is closed");
            }

            //DO NOT BLOCK for isReady. This method must not use requireReady or sendStdin. If this
            //message is buffered in gRPC, NuProcess would have buffered too if we were running the
            //same process locally
            callOnNext(ByteString.copyFrom(buffer), false);
        }

        /**
         * Offers a tombstone task ({@code this}) to the task queue to allow the {@link #run} loop to complete.
         */
        void onComplete() {
            tasks.offer(this);

            //In case the run() loop is blocked waiting for
            synchronized (stdin) {
                stdin.notifyAll();
            }
        }

        /**
         * Calls {@link StreamObserver#onNext} with a fragment containing the provided {@code stdin} data.
         * <p>
         * This method does not check {@link ClientCallStreamObserver#isReady}, and may buffer the message
         * inside gRPC if the observer is not ready.
         *
         * @param data   the {@code stdin} to write
         * @param closed {@code true} if this will be the final chunk of {@code stdin}; otherwise, {@code false}
         */
        private void callOnNext(ByteString data, boolean closed) {
            requestStream.onNext(fragmentCreator.forStdin(RpcStreamChunk.newBuilder()
                    .setClosed(closed)
                    .setData(data)
                    .build()));
        }

        /**
         * Notifies a thread blocked on the {@link #stdin} buffer that the {@link StreamObserver} is ready for
         * another message. If no thread is blocked, this callback does nothing.
         * <p>
         * <b>Note</b>: This callback will be invoked on a gRPC thread.
         */
        private void onReady() {
            synchronized (stdin) {
                //Spurious wakeups will be handled by the loop in requireReady
                stdin.notify();
            }
        }

        /**
         * Waits until the observer is {@link ClientCallStreamObserver#isReady ready}. If the observer is not
         * ready when this method is called, the thread will block on {@link #stdin} until it's notified by
         * either {@link #onComplete} or {@link #onReady}.
         * <p>
         * If the request is {@link #isCanceled canceled} or {@link #isDone completed}, or {@link #closeSent
         * stdin is closed} the observer will never become ready again. If that happens, this method returns
         * {@code false}.
         * <p>
         * <b>Note</b>: If the request's deadline expires the request is canceled. That means checking for
         * cancellation here will also pick up request-level timeouts.
         *
         * @return {@code true} if the observer was ready, or has become ready; otherwise, {@code false} if the
         *         observer has been canceled or {@code stdin} has been closed
         */
        private boolean requireReady() {
            if (requestStream.isReady()) {
                return true;
            }

            synchronized (stdin) {
                while (true) {
                    if (isCanceled() || isDone() || closeSent.get()) {
                        //isCanceled(): _Client side_ cancellation using NioProcess.cancel
                        //isDone(): _Server side_ completion using onCompleted or onError
                        //closeSent.get(): Final stdin RpcStreamChunk written with closed=true
                        return false;
                    }
                    if (requestStream.isReady()) {
                        return true;
                    }

                    try {
                        //The delay here is a backstop, to guard against any edge case where we don't get
                        //notified. It's intentionally high enough that it shouldn't normally get hit
                        stdin.wait(3000L);
                    } catch (InterruptedException e) {
                        if (isCanceled()) {
                            log.debug("Canceled while waiting to write stdin");
                            return false;
                        }
                        Thread.currentThread().interrupt();

                        throw new AbortedException("Interrupted while waiting to write stdin", e);
                    }
                }
            }
        }

        /**
         * Calls {@link ProcessStdioHandler#onStdinReady} to fill the {@link #stdin} buffer and {@link #sendStdin sends}
         * the buffered data to the server. If {@code onStdinReady} returns {@code true}, multiple buffers may be
         * sent before the loop returns.
         * <p>
         * <b>Note</b>: This loop should only be called by the {@link #run} thread, using the {@link #tasks} queue,
         * to ensure multiple threads don't try to call {@code onStdinReady} concurrently. Additionally, this loop
         * may block pending an {@link #onReady} callback from gRPC (potentially multiple times) depending on how
         * many {@code stdin} buffers are filled/sent.
         */
        private void runStdinLoop() {
            if (!writeRequested.compareAndSet(true, false) || closeRequested.get() || isDone()) {
                //If multiple runStdinLoop tasks are queued, or closeStdin is called while the
                //task is queued, or the request (somehow) completes, skip reading stdin
                return;
            }

            boolean more;
            do {
                //Flow: At this point the stdin buffer is in _write mode_
                more = stdioHandler.onStdinReady(stdin);

                //Flow: At this point the stdin buffer is in _read mode_
                boolean closed = closeRequested.get();
                boolean hasStdin = stdin.hasRemaining();
                if (hasStdin || closed) {
                    sendStdin(hasStdin ? ByteString.copyFrom(stdin) : ByteString.EMPTY, closed);
                    if (closed) {
                        closeSent.set(true);
                        requestStream.onCompleted();

                        //Ignore onStdinReady's return value if stdin is closed. Note that
                        //the stdin buffer is left in _read mode_ here because it won't be
                        //used again--subsequent runStdinLoop calls will always no-op
                        break;
                    }
                    if (isDone()) {
                        //If the request (somehow) completes while we're writing stdin,
                        //ignore onStdinReady's return value and break out without calling
                        //onCompleted (because stdin did not complete)
                        break;
                    }
                }

                //Reset the stdin buffer to prepare for the next read. We do this even if
                //onStdinReady wrote nothing to the buffer because we need to ensure the
                //buffer is back in _write mode_ before the next iteration of the loop
                stdin.clear();
                //Flow: At this point the stdin buffer is in _write mode_
            } while (more);
        }

        /**
         * Sends the provided {@code stdin} data (which may be {@link ByteString#EMPTY empty}) and {@code closed}
         * flag using {@link #callOnNext}.
         * <p>
         * <b>Note</b>: If the observer is not {@link ClientCallStreamObserver#isReady ready}, this method will
         * block until it is to prevent buffering messages in gRPC.
         *
         * @param data   the {@code stdin} data to send, which may be empty but not {@code null}
         * @param closed {@code true} if this will be the final {@code stdin} data; otherwise, {@code false}
         */
        private void sendStdin(ByteString data, boolean closed) {
            if (requireReady()) {
                callOnNext(data, closed);
            }
        }
    }
}
