package com.atlassian.bitbucket.mesh.git;

import com.atlassian.bitbucket.mesh.rpc.v1.RpcStreamFragment;
import com.google.common.base.Throwables;
import com.google.common.collect.ImmutableList;
import io.grpc.stub.StreamObserver;

import javax.annotation.Nonnull;
import javax.annotation.Nullable;
import java.util.List;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;

class StreamFragmentResponseObserver implements StreamObserver<RpcStreamFragment> {

    private final ImmutableList.Builder<RpcStreamFragment> fragments = ImmutableList.builder();
    private final CompletableFuture<Void> future = new CompletableFuture<>();

    private boolean completed;
    private Throwable throwable;

    public void await(long timeout, @Nonnull TimeUnit timeUnit) throws InterruptedException, TimeoutException {
        try {
            future.get(timeout, timeUnit);
        } catch (ExecutionException executionException) {
            Throwable cause = executionException.getCause();
            Throwables.throwIfUnchecked(cause);
            throw new RuntimeException(cause);
        }
    }

    @Nonnull
    public List<RpcStreamFragment> getFragments() {
        return fragments.build();
    }

    @Nullable
    public Throwable getThrowable() {
        return throwable;
    }

    public boolean isCompleted() {
        return completed;
    }

    @Override
    public void onNext(RpcStreamFragment value) {
        fragments.add(value);
    }

    @Override
    public void onError(Throwable t) {
        throwable = t;
        future.completeExceptionally(t);
    }

    @Override
    public void onCompleted() {
        completed = true;
        future.complete(null);
    }
}
