package com.atlassian.bitbucket.mesh.grpc;

import com.atlassian.bitbucket.mesh.AbstractStatusException;
import io.grpc.Status;
import io.grpc.stub.StreamObserver;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.annotation.Nonnull;

import static java.util.Objects.requireNonNull;

/**
 * A {@link StreamObserver} implementation which supports <i>bidirectional</i> streaming.
 * <p>
 * For bidirectional streaming, gRPC requires the service to return a {@link StreamObserver} which can receive
 * multiple messages via {@link #onNext}.
 *
 * @param <ReqT> the request message type
 * @param <RespT> the response message type
 */
public class ConversationalRequestObserver<ReqT, RespT> implements StreamObserver<ReqT> {

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

    private final StreamObserver<RespT> responseObserver;
    private final RequestMessageHandler<ReqT, RespT> messageHandler;

    private volatile boolean completed;

    public ConversationalRequestObserver(@Nonnull StreamObserver<RespT> responseObserver,
                                         @Nonnull RequestMessageHandler<ReqT, RespT> messageHandler) {
        this.responseObserver = requireNonNull(responseObserver, "responseObserver");
        this.messageHandler = requireNonNull(messageHandler, "messageHandler");
    }

    @Override
    public void onCompleted() {
        //The only way this observer knows processing is complete is when the client tells it, so when
        //we receive the onCompleted() callback we need to complete the server-side observer to prevent
        //the RPC from lingering until its deadline
        maybeComplete();
    }

    @Override
    public void onError(Throwable t) {
        if (Status.fromThrowable(t).getCode() == Status.Code.CANCELLED) {
            log.debug("The call was canceled");
        } else {
            log.error("The call failed", t);
        }

        //If we get here from our own onError callback in onNext below, this won't do anything. But if
        //we get here because the client canceled the request, or had some other failure, then we need
        //to complete the server-side observer to prevent the RPC from lingering until its deadline
        maybeComplete();
    }

    @Override
    public void onNext(ReqT message) {
        try {
            responseObserver.onNext(messageHandler.onMessage(message));
        } catch (RuntimeException e) {
            if (e instanceof AbstractStatusException) {
                e = ((AbstractStatusException) e).toStatusException();
            }

            completed = true; //maybeComplete() doesn't need to complete the RPC
            responseObserver.onError(e);
        }
    }

    private void maybeComplete() {
        if (!completed) {
            responseObserver.onCompleted();
        }
    }
}
