package com.atlassian.bitbucket.mesh.grpc;

import io.grpc.Context;
import io.grpc.stub.ClientCallStreamObserver;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.annotation.Nullable;

/**
 * A wrapper {@link ClientCallStreamObserver} implementation that limits the speed at which messages can be sent
 * to the server.
 * <p>
 * <em>IMPORTANT</em>: when using {@link BackoffClientCallStreamObserver} it is <em>critical</em> to call
 * {@link #onClosed()} if the RPC is closed (e.g. when {@link io.grpc.stub.StreamObserver#onError} is
 * called on the 'other' {@link io.grpc.stub.StreamObserver}. This is important because {@link ClientCallStreamObserver}
 * does not expose the call state and gRPC does not manage a {@link Context} on the client side.
 * <p>
 * Failure to do so may cause {@link #onNext} to stall indefinitely if the RPC server encounters an error, for example
 * when the server is offline.
 * <p>
 * If the 'other' {@link io.grpc.stub.StreamObserver} extends {@link AbstractFutureResponseObserver},
 * {@link #onClosed()} will already be called when the RPC terminates and no special action is needed.
 *
 * @param <ReqT> the request message type
 */
public class BackoffClientCallStreamObserver<ReqT>
        extends ClientCallStreamObserver<ReqT>
        implements ClientCallCloseListener {

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

    private final ClientCallStreamObserver<ReqT> delegate;
    private final Object monitor;

    private volatile boolean closed;
    private volatile boolean inactive; //inactive rather than active to simplify checks
    private Runnable readyHandler;

    public BackoffClientCallStreamObserver(ClientCallStreamObserver<ReqT> delegate) {
        this(delegate, true);
    }

    public BackoffClientCallStreamObserver(ClientCallStreamObserver<ReqT> delegate, boolean active) {
        this.delegate = delegate;

        inactive = !active;
        monitor = new Object();

        delegate.setOnReadyHandler(this::onReady);
    }

    /**
     * Activates backpressure. <i>Subsequent</i> {@link #onNext messages} will require the observer to be
     * {@link #requireReady ready} before they are sent.
     */
    public void activate() {
        synchronized (monitor) {
            inactive = false;
        }
    }

    @Override
    public void cancel(@Nullable String message, @Nullable Throwable cause) {
        delegate.cancel(message, cause);
    }

    /**
     * Deactivates backpressure. Subsequent {@link #onNext messages} will be sent, or buffered if the observer
     * is not {@link ClientCallStreamObserver#isReady ready}. If a thread is already blocked waiting for the
     * observer to be ready, it is released and allowed to send immediately.
     */
    public void deactivate() {
        synchronized (monitor) {
            inactive = true;
            monitor.notifyAll();
        }
    }

    @Override
    public void disableAutoInboundFlowControl() {
        delegate.disableAutoInboundFlowControl();
    }

    @Override
    public void disableAutoRequestWithInitial(int request) {
        delegate.disableAutoRequestWithInitial(request);
    }

    @Override
    public boolean isReady() {
        return delegate.isReady();
    }

    @Override
    public void onClosed() {
        synchronized (monitor) {
            closed = true;
            monitor.notifyAll();
        }
    }

    @Override
    public void onCompleted() {
        delegate.onCompleted();
    }

    @Override
    public void onError(Throwable t) {
        onClosed();
        delegate.onError(t);
    }

    @Override
    public void onNext(ReqT value) {
        requireReady();

        //If the delegate has been canceled or is already completed, calling onNext here will throw.
        //We intentionally don't try to prevent that. All this observer does is apply backpressure
        delegate.onNext(value);
    }

    @Override
    public void request(int count) {
        delegate.request(count);
    }

    @Override
    public void setMessageCompression(boolean enable) {
        delegate.setMessageCompression(enable);
    }

    @Override
    public void setOnReadyHandler(Runnable onReadyHandler) {
        readyHandler = onReadyHandler;
    }

    private void onReady() {
        synchronized (monitor) {
            monitor.notifyAll();
        }

        if (readyHandler != null) {
            readyHandler.run();
        }
    }

    private void requireReady() {
        //Do leading deactivation/cancellation/ready checks to avoid locking entirely if possible
        if (inactive || closed || delegate.isReady()) {
            return;
        }

        //Only track waiting if we're going to log it. This includes time spent acquiring the lock
        long start = log.isDebugEnabled() ? System.nanoTime() : 0L;
        try {
            synchronized (monitor) {
                while (true) {
                    if (inactive || closed || delegate.isReady()) {
                        if (start != 0L && log.isDebugEnabled()) {
                            long timeMs = (System.nanoTime() - start) / 1_000_000L;
                            if (timeMs > 10L) {
                                log.debug("Waited {}ms for the client to be ready", timeMs);
                            }
                        }
                        return;
                    }

                    //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
                    monitor.wait(3000L);
                }
            }
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
        }
    }
}
