package com.atlassian.bitbucket.mesh.client;

import com.atlassian.bitbucket.mesh.coordination.LogicalClock;
import com.atlassian.bitbucket.mesh.grpc.AbstractFutureResponseObserver;
import com.atlassian.bitbucket.mesh.network.MeshNode;
import com.atlassian.bitbucket.mesh.rpc.v1.*;
import com.atlassian.bitbucket.mesh.transaction.ReplicaConsistencyVerifier;
import com.atlassian.bitbucket.mesh.transaction.ReplicaTransactionCallback;
import com.atlassian.bitbucket.mesh.transaction.ReplicaVersion;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.annotation.Nonnull;
import javax.annotation.Nullable;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CompletionStage;

import static java.util.Objects.requireNonNull;

/**
 * {@link io.grpc.stub.StreamObserver} used on the leader side of a {@code TransactionService/VerifyConsistency}
 * call that handles all interactions with one of the replicas that participate in the consistency verification.
 *
 * @since 1.2
 */
public class ReplicaConsistencyResponseObserver
        extends AbstractFutureResponseObserver<RpcVerifyConsistencyLeaderFragment, RpcVerifyConsistencyFollowerFragment, Void>
        implements ReplicaConsistencyVerifier {

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

    private final ReplicaTransactionCallback callback;
    private final CompletableFuture<Void> future;
    private final LogicalClock logicalClock;
    private final MeshNode replica;
    private final String transactionId;

    ReplicaConsistencyResponseObserver(@Nonnull LogicalClock logicalClock, @Nonnull MeshNode replica,
                                       @Nonnull ReplicateParameters parameters) {
        this.logicalClock = requireNonNull(logicalClock, "logicalClock");
        this.replica = requireNonNull(replica, "replica");

        callback = requireNonNull(parameters, "parameters").getCallback();
        future = new CompletableFuture<>();
        transactionId = parameters.getTransactionId();

        //Set up a listener on the overall request to ensure our stage futures are always completed
        //if the surrounding RPC completes--regardless of how
        addListener(this::wrapUp);
    }

    @Override
    public void cancel(@Nullable String message) {
        //Cancel the RPC; if the RPC isn't complete yet this will trigger wrapUp through the listener we set up
        //in the constructor.
        super.cancel(message);
    }

    @Nonnull
    @Override
    public MeshNode getReplica() {
        return replica;
    }

    @Override
    public void onNext(RpcVerifyConsistencyFollowerFragment fragment) {
        logicalClock.onReceived(fragment.getLogicalTimestamp());
        switch (fragment.getFragmentOneofCase()) {
            case TOKEN_GRANTED:
                onTokenGranted(fragment.getTokenGranted());
                break;
            case YIELD_TOKEN_REQUEST:
                onYieldTokenRequest();
                break;
            default:
                log.warn("[{}] Unexpected follower fragment: {}", this, fragment);
        }
    }

    @Nonnull
    @Override
    public CompletionStage<Void> requestTokenAndReplicaVersion(@Nonnull String repositoryId, long logicalTimestamp) {
        requestStream.onNext(RpcVerifyConsistencyLeaderFragment.newBuilder()
                .setLogicalTimestamp(logicalTimestamp)
                .setTokenRequest(RpcWriteTokenRequest.newBuilder()
                        .setRepository(repositoryId)
                        .setTransactionId(transactionId))
                .build());

        return future;
    }

    @Override
    public void summarize(@Nonnull RpcRepositoryConsistencySummary summary, long logicalTimestamp) {
        requestStream.onNext(RpcVerifyConsistencyLeaderFragment.newBuilder()
                .setLogicalTimestamp(logicalTimestamp)
                .setConsistencySummary(summary)
                .build());
        requestStream.onCompleted();
    }

    private void onTokenGranted(RpcWriteTokenGranted tokenGranted) {
        log.debug("[{}] Granted write token", this);
        ReplicaVersion replicaVersion = ReplicaVersion.fromReplicaVersion(tokenGranted.getReplicaVersion());
        future.complete(null);
        if (!callback.onTokenGranted(replica, replicaVersion)) {
            log.info("[{}] Verification leader rejected the write token granted by {}. " +
                    "This may happen if the verification is in the process of being aborted.", this, replica);
        }
    }

    private void onYieldTokenRequest() {
        if (callback.yieldToken(replica)) {
            log.debug("[{}] Yielded write token back to {} (as requested)", this, replica);
            requestStream.onNext(RpcVerifyConsistencyLeaderFragment.newBuilder()
                    .setLogicalTimestamp(logicalClock.tick())
                    .setTokenYielded(RpcWriteTokenYielded.getDefaultInstance())
                    .build());
        } else {
            log.debug("[{}] Follower {} requested the write token to be yielded, but the verification leader refused " +
                    "because a write lock has already been acquired.", this, replica);
        }
    }

    private void wrapUp(Void result, Throwable thrown) {
        //If we've already received the write token + version info, this will no-op.
        if (thrown == null) {
            future.complete(null);
        } else {
            callback.onReplicaError(replica);
            future.completeExceptionally(thrown);
        }
    }
}
