package com.atlassian.bitbucket.mesh;

import com.google.common.base.Throwables;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.Lists;
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;

public class MockResponseObserver<T> implements StreamObserver<T> {

    private final List<T> fragments = Lists.newCopyOnWriteArrayList();
    private final CompletableFuture<Void> future = new CompletableFuture<>();
    private final Object lock = new Object();

    private boolean completed;
    private Throwable throwable;

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

    public T awaitFragment(int index, long timeout, TimeUnit timeUnit) throws TimeoutException, InterruptedException {
        if (index < fragments.size()) {
            return fragments.get(index);
        }
        long deadline = System.currentTimeMillis() + timeUnit.toMillis(timeout);
        synchronized (lock) {
            while (!completed && fragments.size() <= index && System.currentTimeMillis() < deadline) {
                lock.wait(Math.max(0L, deadline - System.currentTimeMillis()));
            }
            if (index < fragments.size()) {
                return fragments.get(index);
            }
            throw new TimeoutException("Timed out waiting for fragment (completed = " + completed + ")");
        }
    }

    @Nonnull
    public List<T> getFragments() {
        return ImmutableList.copyOf(fragments);
    }

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

    public boolean isCompleted() {
        return completed;
    }

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

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

    @Override
    public void onNext(T value) {
        synchronized (lock) {
            fragments.add(value);
            lock.notify();
        }
    }
}