package com.atlassian.bitbucket.mesh.transaction;

import com.atlassian.bitbucket.mesh.repository.ManagedRepository;
import com.atlassian.bitbucket.mesh.rpc.v1.*;
import com.google.common.base.Throwables;
import io.grpc.stub.ClientCallStreamObserver;
import io.grpc.stub.ClientResponseObserver;
import org.junit.Assert;

import javax.annotation.Nonnull;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;
import java.util.function.Consumer;

class BlockingReplicaTransaction
        implements ClientResponseObserver<RpcTransactLeaderFragment, RpcTransactFollowerFragment> {

    private final CompletableFuture<RpcCommitTransactionResponse> commitFuture;
    private final CompletableFuture<RpcPrepareTransactionResponse> prepareFuture;
    private final ManagedRepository repository;
    private final CompletableFuture<RpcRollBackTransactionResponse> rollBackFuture;

    private volatile ClientCallStreamObserver<RpcTransactLeaderFragment> requestStream;

    BlockingReplicaTransaction(ManagedRepository repository) {
        this.repository = repository;

        commitFuture = new CompletableFuture<>();
        prepareFuture = new CompletableFuture<>();
        rollBackFuture = new CompletableFuture<>();
    }

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

    public void cancel(String message, Throwable error) {
        requestStream.cancel(message, error);
    }

    public void commit() {
        sendRequestAndWaitForReply(builder -> builder.setCommitRequest(RpcCommitTransactionRequest.newBuilder()
                        .setVersion(repository.getVersion())),
                commitFuture);
    }

    @Nonnull
    public RpcPrepareTransactionResponse prepare(@Nonnull RpcTransaction operation) {
        return sendRequestAndWaitForReply(builder -> builder.setPrepareRequest(
                RpcPrepareTransactionRequest.newBuilder()
                        .setLeader(AbstractPreparedActionWriteTest.LEADER_ID)
                        .setTimeout(AbstractPreparedActionWriteTest.TX_TIMEOUT_MS)
                        .setTransaction(operation)), prepareFuture);
    }

    @Override
    public void onError(Throwable t) {
        commitFuture.completeExceptionally(t);
        prepareFuture.completeExceptionally(t);
        rollBackFuture.completeExceptionally(t);
    }

    @Override
    public void onNext(RpcTransactFollowerFragment value) {
        switch (value.getFragmentOneofCase()) {
            case COMMIT_RESPONSE:
                commitFuture.complete(value.getCommitResponse());
                break;
            case PREPARE_RESPONSE:
                prepareFuture.complete(value.getPrepareResponse());
                break;
            case ROLL_BACK_RESPONSE:
                rollBackFuture.complete(value.getRollBackResponse());
                break;
            default:
                Assert.fail("Unexpected fragment type");
        }
    }

    @Override
    public void onCompleted() {
    }

    public void rollBack() {
        sendRequestAndWaitForReply(builder -> builder.setRollBackRequest(
                RpcRollBackTransactionRequest.getDefaultInstance()), rollBackFuture);
    }

    public void summarize(boolean committed, String contentHash) {
        requestStream.onNext(RpcTransactLeaderFragment.newBuilder()
                .setSummary(RpcTransactionSummary.newBuilder()
                        .setCommitted(committed)
                        .setContentHash(contentHash))
                .build());
        requestStream.onCompleted();
    }

    boolean isStarted() {
        return requestStream != null;
    }

    private <T> T sendRequestAndWaitForReply(Consumer<RpcTransactLeaderFragment.Builder> requestBuilder,
                                             CompletableFuture<T> future) {
        RpcTransactLeaderFragment.Builder builder = RpcTransactLeaderFragment.newBuilder();
        requestBuilder.accept(builder);
        requestStream.onNext(builder.build());

        try {
            return future.get(AbstractPreparedActionWriteTest.TX_TIMEOUT_MS, TimeUnit.MILLISECONDS);
        } catch (ExecutionException e) {
            Throwable cause = e.getCause();
            Throwables.throwIfUnchecked(cause);
            throw new RuntimeException("Error", e);
        } catch (InterruptedException e) {
            requestStream.cancel("Interrupted", e);
            Thread.currentThread().interrupt();
            throw new RuntimeException("Interrupted", e);
        } catch (TimeoutException e) {
            requestStream.cancel("Timed out", e);
            throw new RuntimeException("Timed out", e);
        }
    }
}
