package com.atlassian.bitbucket.mesh.grpc;

import com.google.protobuf.MessageLite;
import io.grpc.stub.ServerCallStreamObserver;
import io.grpc.stub.StreamObserver;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import static com.google.common.base.Preconditions.checkArgument;

public class BackoffStreamObserver<T extends MessageLite> extends ServerCallStreamObserver<T> {

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

    private final ServerCallStreamObserver<T> delegate;
    private final Object lock;

    private volatile Runnable delegateCancelHandler;

    public BackoffStreamObserver(StreamObserver<T> delegate) {
        checkArgument(delegate instanceof ServerCallStreamObserver, "must be a ServerCallStreamObserver");

        this.delegate = (ServerCallStreamObserver<T>) delegate;
        this.delegate.setOnCancelHandler(new OnCancelHandler());
        this.delegate.setOnReadyHandler(new OnReadyHandler());

        lock = new Object();
    }

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

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

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

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

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

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

    @Override
    public void onNext(T 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 setCompression(String compression) {
        delegate.setCompression(compression);
    }

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

    @Override
    public void setOnCancelHandler(Runnable onCancelHandler) {
        delegateCancelHandler = onCancelHandler;
    }

    @Override
    public void setOnReadyHandler(Runnable onReadyHandler) {
        throw new IllegalStateException("An OnReadyHandler has already been set to apply back-pressure");
    }

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

    private void requireReady() {
        //Do leading cancellation/ready checks to avoid locking entirely if possible
        if (delegate.isCancelled() || 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;
        synchronized (lock) {
            while (true) {
                if (delegate.isCancelled() || 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;
                }

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

    private class OnCancelHandler implements Runnable {

        @Override
        public void run() {
            synchronized (lock) {
                log.trace("Client canceled");
                lock.notifyAll();
            }
            Runnable delegateHandler = delegateCancelHandler;
            if (delegateHandler != null) {
                delegateHandler.run();
            }
        }
    }

    private class OnReadyHandler implements Runnable {

        @Override
        public void run() {
            synchronized (lock) {
                log.trace("Client is ready for more messages");
                lock.notifyAll();
            }
        }
    }
}
