package com.atlassian.bitbucket.mesh.git;

import com.atlassian.bitbucket.mesh.git.exception.AbortedException;
import com.google.common.base.Throwables;
import io.grpc.stub.ClientCallStreamObserver;
import io.grpc.stub.ClientResponseObserver;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.concurrent.*;

class UploadResponseObserver<ReqT, RespT> implements ClientResponseObserver<ReqT, RespT> {

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

    private final CountDownLatch completed = new CountDownLatch(1);
    private final Object monitor = new Object();

    private volatile ClientCallStreamObserver<ReqT> requestStream;
    private volatile CompletableFuture<RespT> responseFuture = new CompletableFuture<>();

    @Override
    public void beforeStart(ClientCallStreamObserver<ReqT> requestStream) {
        this.requestStream = requestStream;

        requestStream.setOnReadyHandler(() -> {
            synchronized (monitor) {
                monitor.notifyAll();
            }
        });
    }

    Future<RespT> sendRequest(ReqT request) {
        requestStream.onNext(request);
        return responseFuture;
    }

    void sendChunk(ReqT fragment) {
        if (awaitReady()) {
            requestStream.onNext(fragment);
        }
    }

    @Override
    public void onNext(RespT value) {
        responseFuture.complete(value);
    }

    @Override
    public void onError(Throwable t) {
        log.error("Received an error!", t);
        if (!responseFuture.isDone()) {
            responseFuture.completeExceptionally(t);
        }
        completed.countDown();
    }

    @Override
    public void onCompleted() {
        completed.countDown();
    }

    void awaitCompleted(long timeout, TimeUnit timeUnit) throws InterruptedException {
        completed.await(timeout, timeUnit);
    }

    boolean awaitReady() {
        if (requestStream.isReady()) {
            return true;
        }

        while (true) {
            try {
                synchronized (monitor) {
                    if (completed.getCount() == 0 || requestStream.isReady()) {
                        return true;
                    }

                    monitor.wait(3_000);
                }
            } catch (InterruptedException e) {
                if (responseFuture.isCancelled()) {
                    log.debug("Canceled while writing object content");
                    return false;
                }
                Thread.currentThread().interrupt();

                throw new AbortedException("Interrupted while writing object content", e);
            }
        }
    }

    RespT awaitResponse(long timeout, TimeUnit timeUnit) throws TimeoutException {
        try {
            RespT response = responseFuture.get(timeout, timeUnit);

            //Prepare for the next response
            responseFuture = new CompletableFuture<>();
            return response;
        } 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("TODO: Handle unexpected errors better", e);
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();

            throw new RuntimeException("TODO: Handle interrupts better", e);
        }
    }
}
