package com.atlassian.bitbucket.mesh.git;

import com.atlassian.bitbucket.mesh.AbstractStatusException;
import com.atlassian.bitbucket.mesh.git.exception.CommandTimeoutException;
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.process.handler.AbstractProcessStdioHandler;
import com.atlassian.bitbucket.mesh.rpc.util.ByteStringUtils;
import com.atlassian.bitbucket.mesh.rpc.v1.RpcStreamChunk;
import com.atlassian.bitbucket.mesh.rpc.v1.RpcStreamFragment;
import io.grpc.Context;
import io.grpc.stub.ServerCallStreamObserver;
import io.grpc.stub.StreamObserver;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.annotation.Nonnull;
import java.nio.ByteBuffer;
import java.util.concurrent.ForkJoinPool;

import static java.util.Objects.requireNonNull;

/**
 * A {@link ProcessStdioHandler ProcessStdioHandler} implementation which streams {@code stderr}, {@code stdout}
 * and the exit code using {@link RpcStreamFragment}s.
 * <p>
 * Note: This handler doesn't support providing {@code stdin}. It's intended solely for use streaming fragments
 * of output.
 */
public class FragmentProcessStdioHandler<T>
        extends AbstractProcessStdioHandler<Void>
        implements ProcessCompletionHandler {

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

    private final Context context;
    private final RpcResponseFragmentFactory<T> fragmentFactory;
    private final ServerCallStreamObserver<T> observer;
    private final Object readyLock;

    public FragmentProcessStdioHandler(@Nonnull StreamObserver<T> o,
                                       @Nonnull RpcResponseFragmentFactory<T> fragmentFactory) {
        if (o instanceof ServerCallStreamObserver) {
            observer = (ServerCallStreamObserver<T>) o;
        } else {
            requireNonNull(o, "observer");

            throw new IllegalArgumentException("A ServerCallStreamObserver is required");
        }

        this.fragmentFactory = requireNonNull(fragmentFactory, "fragmentFactory");

        context = Context.current();
        //Wake up any threads blocked in requireReady to have them exit early. If for some reason context is not
        //cancellable, those threads will eventually discover that the context has closed and unblock (after max 3s).
        context.addListener(ctx -> onReady(), ForkJoinPool.commonPool());

        //Create the lock before setting the handlers to ensure it can't be null when they're called
        readyLock = new Object();

        observer.setOnCancelHandler(this::onCancel);
        observer.setOnReadyHandler(this::onReady);
    }

    @Override
    public Void getResult() {
        return null;
    }

    @Override
    public void onComplete(@Nonnull ProcessSummary summary) {
        Throwable thrown = summary.getThrown();
        if (thrown == null && summary.isTimedOut()) {
            thrown = new CommandTimeoutException(summary);
        }

        if (summary.isCanceled() || thrown == null) {
            observer.onCompleted();
        } else {
            observer.onError(thrown instanceof AbstractStatusException ?
                    ((AbstractStatusException) thrown).toStatusException() : thrown);
        }
    }

    @Override
    public void onExit(int exitCode) {
        //The cancellation handling in callOnNext is safe here because the process code checks to ensure
        //the process is still running before trying to cancel it, so cancel() will be a harmless no-op
        callOnNext(fragmentFactory.createResultFragment(exitCode));
    }

    @Override
    public void onStart(@Nonnull OsProcess process) {
        callOnNext(fragmentFactory.createStartFragment(process));
    }

    @Override
    public void onStderr(@Nonnull ByteBuffer buffer, boolean closed) {
        callOnNext(fragmentFactory.createStderrFragment(toChunk(buffer, closed)));
    }

    @Override
    public boolean onStdinReady(@Nonnull ByteBuffer buffer) {
        buffer.flip();
        closeStdin();

        return false;
    }

    @Override
    public void onStdout(@Nonnull ByteBuffer buffer, boolean closed) {
        callOnNext(fragmentFactory.createStdoutFragment(toChunk(buffer, closed)));
    }

    /**
     * Builds a {@link RpcStreamChunk} from the provided {@link ByteBuffer} and {@code closed} flag.
     * <p>
     * The implementation here copies from the provided buffer rather than trying to wrap it because NuProcess
     * reuses a single buffer for every stderr callback and another for every stdout callback. Wrapping would
     * mean those contents could be replaced by a subsequent callback before gRPC sent them.
     *
     * @param buffer the buffer to copy into the new chunk
     * @param closed {@code true} if the stream has been closed; otherwise {@code false}
     * @return the built chunk
     */
    protected static RpcStreamChunk toChunk(ByteBuffer buffer, boolean closed) {
        return RpcStreamChunk.newBuilder()
                .setClosed(closed)
                .setData(ByteStringUtils.unsafeWrap(buffer))
                .build();
    }

    /**
     * @return {@code true} if the {@link OsProcess}, the {@link StreamObserver} has been canceled or the gRPC context
     *         has closed; otherwise, {@code false}
     */
    @Override
    protected boolean isCanceled() {
        return super.isCanceled() || observer.isCancelled() || context.isCancelled();
    }

    /**
     * Invokes {@link StreamObserver#onNext} with the provided {@link RpcStreamFragment} when the observer
     * {@link #requireReady() is ready}.
     * <p>
     * Some processes (especially hosting processes like {@code git http-backend} and {@code git upload-pack})
     * are likely to be able to produce output significantly faster than that output can be relayed via gRPC.
     * Blindly calling {@code onNext} can result in buffering a large number of fragments in memory, which can
     * eventually exhaust the heap. <i>gRPC itself applies no backpressure for {@code onNext};</i> messages
     * are appended to an <i>unbounded</i> queue if the client hasn't requested them.
     *
     * @param fragment the fragment to send
     */
    private void callOnNext(T fragment) {
        if (requireReady()) {
            observer.onNext(fragment);
        }
    }

    /**
     * Cancels the process and then notifies the {@link #readyLock} in case the process thread is blocked
     * in {@link #callOnNext}.
     */
    private void onCancel() {
        cancel();

        //In case the process thread is blocked waiting for queue space when the request is canceled,
        //acquire the readyLock and signal all waiters (which should detect cancellation and abort)
        synchronized (readyLock) {
            readyLock.notifyAll(); //If no threads are blocked, this will no-op
        }
    }

    /**
     * Notifies the {@link #readyLock} to release the process thread if it's blocked in {@link #callOnNext}.
     * If the process thread is not blocked, notifying does nothing.
     */
    private void onReady() {
        synchronized (readyLock) {
            readyLock.notify(); //If no threads are blocked, this will no-op
        }
    }

    /**
     * Waits until the observer is {@link ServerCallStreamObserver#isReady ready}. If the observer is not ready
     * when this method is called, the thread will block on {@link #readyLock} until it's notified by either
     * {@link #onCancel} or {@link #onReady}.
     * <p>
     * If the observer is {@link ServerCallStreamObserver#isCancelled canceled} it will never again become ready.
     * If that happens, this method will return {@code false} to relay that to the caller.
     * <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, indicating the request been terminated
     */
    private boolean requireReady() {
        //Do an up-front isReady() check to try and avoid locking
        if (observer.isReady()) {
            return true;
        }

        synchronized (readyLock) {
            while (true) {
                if (isCanceled()) {
                    //If the observer isn't ready because the request has been canceled, don't try
                    //to wait; we may end up deadlocking. Just drop the fragment instead
                    return false;
                }
                if (observer.isReady()) {
                    return true;
                }

                //The observer still isn't ready, and the request hasn't been canceled, so wait
                //for the ready callback (or cancellation)
                try {
                    //The delay here is a backstop, to guard against any edge case where we don't
                    //get notified. It's high enough that it shouldn't normally get hit
                    readyLock.wait(3000L);
                } catch (InterruptedException e) {
                    if (isCanceled()) {
                        //If the request was canceled while we were waiting, no worries
                        log.debug("Canceled while waiting to send fragment");
                        return false;
                    }
                    Thread.currentThread().interrupt();

                    throw new IllegalStateException("Interrupted while waiting to send fragment");
                }
            }
        }
    }
}
