package com.atlassian.bitbucket.mesh.grpc;

import com.google.common.base.Throwables;
import io.grpc.Status;
import io.grpc.stub.ClientCallStreamObserver;
import io.grpc.stub.ClientResponseObserver;

import javax.annotation.Nonnull;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.Future;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.function.BiConsumer;

/**
 * A base class for constructing {@link ClientResponseObserver}s which use a {@link CompletableFuture} to allow
 * canceling, or waiting for, a request to complete.
 *
 * @param <ReqT>  the request's message type
 * @param <RespT> the response's message type
 * @param <T>     the result type from processing the response (or responses)
 */
public abstract class AbstractFutureResponseObserver<ReqT, RespT, T> implements ClientResponseObserver<ReqT, RespT> {

    protected volatile ClientCallStreamObserver<ReqT> requestStream;

    private final CompletableFuture<T> future;
    private final AtomicBoolean requestCanceled;

    private volatile T result;

    protected AbstractFutureResponseObserver() {
        future = new CompletableFuture<>();
        requestCanceled = new AtomicBoolean();

        addListener((result, exception) -> {
            if (requestStream instanceof ClientCallCloseListener) {
                ((ClientCallCloseListener) requestStream).onClosed();
            }
            if (future.isCancelled() && requestCanceled.compareAndSet(false, true)) {
                ClientCallStreamObserver<ReqT> r = requestStream;
                if (r != null) {
                    r.cancel("The request was canceled by the caller", exception);
                }
            }
        });
    }

    /**
     * Waits until the request completes and returns the result, if {@link #onCompleted} was called, or the exception,
     * if {@link #onError} was called.
     *
     * @return the request's result
     */
    public T asResult() {
        try {
            return future.get();
        } catch (ExecutionException e) {
            Throwable cause = e.getCause();
            Throwables.throwIfUnchecked(cause);

            //This should never actually happen. Even if we don't "understand" the error we
            //got in onError, it should be a StatusRuntimeException and be propagated above
            throw new RuntimeException("Unexpected error", e);
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();

            throw new RuntimeException("Interrupted while waiting for a gRPC response", e);
        }
    }

    /**
     * Stores the {@link ClientCallStreamObserver} so canceling the future can be propagated to the underlying
     * request. If the future has already been canceled, it's propagated to the request here.
     *
     * @param requestStream the request stream, which can be used to cancel or monitor the request
     */
    @Override
    public void beforeStart(ClientCallStreamObserver<ReqT> requestStream) {
        if (future.isCancelled() && requestCanceled.compareAndSet(false, true)) {
            requestStream.cancel("The request was canceled before it could be started", null);
        } else {
            this.requestStream = requestStream;
        }
    }

    /**
     * Applies any {@link #setResult set result} to complete the request's {@link Future}. If no result has
     * been set, <i>the future will be completed with a {@code null} result</i>.
     */
    @Override
    public void onCompleted() {
        future.complete(result);
    }

    /**
     * Uses the provided {@link Throwable} to complete the request's {@link Future}, if it hasn't already been
     * completed (e.g. canceled).
     *
     * @param t the error
     */
    @Override
    public void onError(Throwable t) {
        if (future.isDone()) {
            //If the future has already been completed (local cancellation, most likely), there's no
            //reason to process the exception (we have nowhere to "put" it)
            return;
        }

        //If the exception has a status code attached, and that status is CANCELLED, then propagate
        //that cancellation to the future instead of the exception
        Status status = Status.fromThrowable(t);
        if (status.getCode() == Status.Code.CANCELLED && requestCanceled.compareAndSet(false, true)) {
            //This cancellation could be _local or remote_. The way gRPC has implemented cancellation
            //means the client receives an onError for its own cancellations, not just remote ones
            future.cancel(true); //false or true is irrelevant to CompletableFuture
        } else {
            future.completeExceptionally(t);
        }
    }

    /**
     * Adds a listener which will be called when request processing completes.
     *
     * @param listener the listener to add
     */
    protected void addListener(@Nonnull BiConsumer<? super T, ? super Throwable> listener) {
        future.whenComplete(listener);
    }

    protected void cancel(String message) {
        if (requestStream != null && requestCanceled.compareAndSet(false, true)) {
            requestStream.cancel(message, null);
        }
    }

    /**
     * Checks whether the {@link Future} has been canceled.
     *
     * @return {@code true} if the future has been canceled; otherwise, {@code false}
     */
    protected boolean isCanceled() {
        return future.isCancelled() || requestCanceled.get();
    }

    /**
     * Checks whether the {@link Future} has been completed, exceptionally or with a result. This will only be
     * {@code true} after {@link #onCompleted} or {@link #onError} have been called, to set a result or error,
     * respectively, or if the request has been canceled.
     *
     * @return {@code true} if the future is done; otherwise, {@code false}
     */
    protected boolean isDone() {
        return future.isDone();
    }

    /**
     * Sets the result to be applied to the {@link Future} when {@link #onCompleted} is called. If {@code onCompleted}
     * is not called, the result will be discarded.
     *
     * @param result the result to set
     */
    protected void setResult(T result) {
        this.result = result;
    }
}
