package com.atlassian.bitbucket.mesh.client;

import com.atlassian.bitbucket.mesh.coordination.LogicalClock;
import com.atlassian.bitbucket.mesh.gprc.SynchronizedClientCallStreamObserver;
import com.atlassian.bitbucket.mesh.grpc.AbstractFutureResponseObserver;
import com.atlassian.bitbucket.mesh.grpc.BackoffClientCallStreamObserver;
import com.atlassian.bitbucket.mesh.network.MeshNode;
import com.atlassian.bitbucket.mesh.repository.Repository;
import com.atlassian.bitbucket.mesh.rpc.v1.*;
import com.atlassian.bitbucket.mesh.transaction.*;
import com.atlassian.bitbucket.mesh.util.SatiableConsumer;
import com.google.common.base.Throwables;
import io.grpc.stub.ClientCallStreamObserver;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.annotation.Nonnull;
import java.time.Duration;
import java.util.Optional;
import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicReference;
import java.util.function.UnaryOperator;

import static com.atlassian.bitbucket.mesh.rpc.v1.RpcPrepareTransactionResponse.Vote.ABSTAIN;
import static com.atlassian.bitbucket.mesh.rpc.v1.RpcPrepareTransactionResponse.Vote.YES;
import static com.atlassian.bitbucket.mesh.transaction.TransactionState.*;
import static java.util.Objects.requireNonNull;

class ReplicaTransactionResponseObserver
        extends AbstractFutureResponseObserver<RpcTransactLeaderFragment, RpcTransactFollowerFragment, Void>
        implements ReplicaTransaction {

    private static final UnaryOperator<TransactionState> MAYBE_COMMIT =
            current -> current == PREPARED || current == PREPARING ? COMMITTING : current;
    private static final UnaryOperator<TransactionState> MAYBE_ROLL_BACK =
            current -> current == PREPARED || current == PREPARING ? ROLLING_BACK : current;

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

    private final ReplicaTransactionCallback callback;
    private final AtomicBoolean completed;
    private final CompletableFuture<TransactionState> finalizeStage;
    private final MeshNode leader;
    private final LogicalClock logicalClock;
    private final CompletableFuture<PrepareResult> prepareStage;
    private final MeshNode replica;
    private final AtomicReference<TransactionState> state;
    private final String transactionId;

    private volatile BackoffClientCallStreamObserver<RpcTransactLeaderFragment> backoffObserver;
    private volatile RpcPackWriter packWriter;
    private volatile PrepareResult prepareResult;
    private volatile ReplicaVersion replicaVersion;
    private volatile String transactionType;

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

        callback = parameters.getCallback();
        completed = new AtomicBoolean();
        finalizeStage = new CompletableFuture<>();
        leader = requireNonNull(parameters, "parameters").getLeader();
        prepareStage = new CompletableFuture<>();
        state = new AtomicReference<>(INITIAL);
        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 boolean await(@Nonnull Duration timeout) {
        Throwable cause;
        try {
            return finalizeStage.get(timeout.getSeconds(), TimeUnit.SECONDS) == COMMITTED;
        } catch (ExecutionException e) {
            Throwables.throwIfUnchecked(e.getCause());
            cause = e.getCause();
            //Fall-through intentional
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
            cause = e;
            //Fall-through intentional
        } catch (TimeoutException e) {
            //The RPC isn't canceled here so we can still try to send a summary, to inform the follower
            //it's now considered inconsistent. If the follower actually arrived at the expected outcome,
            //its subsequent repair should effectively no-op
            cause = e;
            //Fall-through intentional
        }

        //We shouldn't get here, but...
        throw chooseFinalizeError(cause);
    }

    @Override
    public void beforeStart(ClientCallStreamObserver<RpcTransactLeaderFragment> requestStream) {
        backoffObserver = new BackoffClientCallStreamObserver<>(
                new SynchronizedClientCallStreamObserver<>(requestStream), false);

        super.beforeStart(backoffObserver);
    }

    @Override
    public void commit(long logicalTimestamp) {
        //This allows two possible transitions:
        //- PREPARING -> COMMITTING
        //- PREPARED -> COMMITTING
        //If a quorum is reached before all followers have voted, the transaction may be committed before
        //some followers have voted. In those cases we transition from PREPARING to COMMITTING. If any of
        //the slower followers vote NO, they'll ignore the optimistic commit request.
        if (state.updateAndGet(MAYBE_COMMIT) == COMMITTING) {
            requestStream.onNext(RpcTransactLeaderFragment.newBuilder()
                    .setCommitRequest(RpcCommitTransactionRequest.getDefaultInstance())
                    .setLogicalTimestamp(logicalTimestamp)
                    .build());
        } else {
            log.debug("[{}] Not sending commit request; current state is {}", this, state);
        }
    }

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

    @Nonnull
    @Override
    public Optional<ReplicaVersion> getReplicaVersion() {
        return Optional.ofNullable(replicaVersion);
    }

    @Nonnull
    @Override
    public Optional<PrepareResult> getPrepareResult() {
        return Optional.ofNullable(prepareResult);
    }

    @Nonnull
    @Override
    public SatiableConsumer<RpcStreamChunk> newPackWriter(@Nonnull Repository repository) {
        if (isDone()) {
            log.debug("[{}] Opting out of pack replication because the RPC has already terminated", this);
            return chunk -> true;
        }
        packWriter = new RpcPackWriter(repository);
        return packWriter;
    }

    @Override
    public void onNext(RpcTransactFollowerFragment fragment) {
        logicalClock.onReceived(fragment.getLogicalTimestamp());
        switch (fragment.getFragmentOneofCase()) {
            case ABORT_STAGING_PACK:
                onAbortStagingPack();
                break;
            case COMMIT_RESPONSE:
                onCommitted(fragment.getCommitResponse());
                break;
            case PREPARE_RESPONSE:
                onPrepared(fragment.getPrepareResponse());
                break;
            case ROLL_BACK_RESPONSE:
                onRolledBack(fragment.getRollBackResponse());
                break;
            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<PrepareResult> prepare(@Nonnull Iterable<RpcTransaction> transactionFragments,
                                                  long logicalTimestamp) {
        requireNonNull(transactionFragments, "transactionFragments");

        TransactionState previousState = state.getAndSet(PREPARING);
        if (previousState != INITIAL) {
            log.warn("[{}] Unexpected prepare from {} state", this, previousState);
        }

        for (RpcTransaction transaction : transactionFragments) {
            transactionType = transaction.getDetailsOneofCase().toString();
            requestStream.onNext(RpcTransactLeaderFragment.newBuilder()
                    .setPrepareRequest(RpcPrepareTransactionRequest.newBuilder()
                            .setLeader(leader.getId())
                            .setTransaction(transaction))
                    .setLogicalTimestamp(logicalTimestamp)
                    .build());
        }
        return prepareStage;
    }

    @Override
    public void rollBack(long logicalTimestamp) {
        //This allows two possible transitions:
        //- PREPARING -> ROLLING_BACK
        //- PREPARED -> ROLLING_BACK
        //In some cases it may be possible to decide that the transaction should roll back before every
        //follower has voted. In those cases we transition from PREPARING to ROLLING_BACK on followers
        //that haven't voted yet
        if (state.updateAndGet(MAYBE_ROLL_BACK) == ROLLING_BACK) {
            requestStream.onNext(RpcTransactLeaderFragment.newBuilder()
                    .setRollBackRequest(RpcRollBackTransactionRequest.getDefaultInstance())
                    .setLogicalTimestamp(logicalTimestamp)
                    .build());
        } else {
            log.debug("[{}] Not sending roll back request; current state is {}", this, state);
        }
    }

    @Override
    public void summarize(@Nonnull RpcTransactionSummary summary, long logicalTimestamp) {
        if (completed.compareAndSet(false, true)) {
            //Ensure we don't deadlock due to unexpected backpressure from staging a pack. If backpressure
            //is already deactivated, this will do nothing
            backoffObserver.deactivate();

            requestStream.onNext(RpcTransactLeaderFragment.newBuilder()
                    .setSummary(summary)
                    .setLogicalTimestamp(logicalTimestamp)
                    .build());
            requestStream.onCompleted();
        } else {
            log.debug("[{}] Discarding summary; the RPC is already complete", this);
        }
    }

    @Override
    public String toString() {
        return transactionId + "@" + replica.getName();
    }

    private RuntimeException chooseFinalizeError(Throwable cause) {
        TransactionState state = this.state.get();
        if (state == COMMITTED || state == COMMITTING) {
            return new TransactionCommitFailedException("Unexpected exception committing prepared transaction",
                    RpcTransactionCommitFailed.Reason.UNKNOWN, cause);
        }

        return new TransactionRollbackFailedException("Unexpected exception rolling back prepared transaction",
                RpcTransactionRollbackFailed.Reason.UNKNOWN, cause);
    }

    private void onAbortStagingPack() {
        RpcPackWriter writer = packWriter;
        if (writer != null) {
            writer.abort();
        }
    }

    private void onCommitted(RpcCommitTransactionResponse response) {
        if (response.hasCommitFailed()) {
            TransactionState previousState = state.getAndSet(ROLLED_BACK);
            log.warn("[{}] Failed to commit: {} (State: {}))",
                    this, response.getCommitFailed().getReason(), previousState);
        } else {
            TransactionState previousState = state.getAndSet(COMMITTED);
            if (previousState != COMMITTING) {
                log.warn("[{}] Unexpected commit from {} state", this, previousState);
            }
            replicaVersion = ReplicaVersion.fromReplicaVersion(response.getReplicaVersion());
        }

        finalizeStage.complete(state.get());
    }

    private void onPrepared(RpcPrepareTransactionResponse response) {
        RpcPrepareTransactionResponse.Vote vote = response.getVote();

        if (vote == YES) {
            prepareResult = PrepareResult.SUCCESS;
            if (response.hasWriteTokenGranted()) {
                onTokenGranted(response.getWriteTokenGranted());
            }

            //If we have optimistically sent a commit request, or we've already concluded the transaction is
            //going to roll back, we may have transitioned "past" PREPARED already. In such cases this won't
            //update the state (and needs to not do so). Otherwise, this will transition to PREPARED
            if (state.compareAndSet(PREPARING, PREPARED)) {
                log.debug("[{}] Prepared remote transaction", this);
            } else {
                log.debug("[{}] Remote preparation completed after transition to {}", this, state);
            }
        } else if (vote == ABSTAIN) {
            //The follower has opted out of voting on the transaction. An example of why a follower might do
            //this is if it was just assigned a new partition replica and has not yet completed backfilling
            //all the necessary repositories. In that case it is unable to prepare any transactions, so it
            //would otherwise have to vote NO. That can be undesirable, though. Consider the case where the
            //number of replicas for a partition has increased from 1 to 2. If the new follower votes NO on
            //every transaction while it backfills, all writes will be blocked because a quorum cannot form
            log.info("[{}] The node has abstained from voting (State: {}; Type: {})", this, state, transactionType);
            prepareResult = PrepareResult.ABSTAIN;

            state.set(ROLLED_BACK);
            finalizeStage.complete(ROLLED_BACK);
        } else {
            //If a follower fails to prepare the transaction it will send a NO vote, but then it still has to
            //wait for the summary to find out whether it's going to become inconsistent or not.
            log.warn("[{}] Failed to prepare remote transaction (Vote: {}; State: {}; Type: {})", this, vote, state,
                    transactionType);
            prepareResult = PrepareResult.FAILURE;

            //From the follower's perspective, the transaction is already rolled back. Transition directly to
            //that state and then, in addition to completing the prepareStage below, complete finalizeStage
            state.set(ROLLED_BACK);
            finalizeStage.complete(ROLLED_BACK);
        }

        prepareStage.complete(prepareResult);
    }

    private void onRolledBack(RpcRollBackTransactionResponse response) {
        if (response.hasRollbackFailed()) {
            TransactionState previousState = state.getAndSet(COMMITTED);
            log.warn("[{}] Failed to roll back: {} (State: {}))",
                    this, response.getRollbackFailed().getReason(), previousState);
        } else {
            TransactionState previousState = state.getAndSet(ROLLED_BACK);
            if (!previousState.isRollbackOnly()) { //Allow both ROLLED_BACK and ROLLING_BACK
                log.warn("[{}] Unexpected rollback from {} state", this, previousState);
            }
        }

        finalizeStage.complete(state.get());
    }

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

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

    private void wrapUp(Void result, Throwable thrown) {
        //If we've already prepared, this will no-op. We don't check because there's complex processing
        //here to decide how to complete the future
        if (thrown == null) {
            prepareStage.complete(PrepareResult.FAILURE);
        } else {
            prepareStage.completeExceptionally(thrown);
        }
        if (finalizeStage.isDone()) {
            //If the finalizeStage is already complete, skip out trying to decide to complete it
            return;
        }

        //If no exception has been thrown and the internal state is COMMITTED or ROLLED_BACK,
        //we can complete the finalizeStage normally.
        TransactionState state = this.state.get();
        if (thrown == null && state.isFinished()) {
            finalizeStage.complete(state);
            return;
        }

        //If finalizeStage hasn't completed, we need to assume it failed. _How_ to assume that will
        //depend a bit on how the RPC completed
        RuntimeException finalizeException = null;
        RpcGitError error = thrown == null ? null : RpcGitErrorUtils.fromThrowable(thrown);
        if (error != null) {
            //We only really expect transaction-related errors here; for anything else we'll
            //end up using a catch-all exception below
            switch (error.getErrorOneofCase()) {
                case ILLEGAL_TRANSACTION_STATE:
                    RpcIllegalTransactionState illegalTransactionState = error.getIllegalTransactionState();
                    finalizeException = new IllegalTransactionStateException(thrown.getMessage(),
                            illegalTransactionState.getTransactionId(),
                            TransactionState.valueOf(illegalTransactionState.getState()),
                            TransactionState.valueOf(illegalTransactionState.getTargetState()));
                    break;
                case TRANSACTION_COMMIT_FAILED:
                    finalizeException = new TransactionCommitFailedException(thrown.getMessage(),
                            error.getTransactionCommitFailed().getReason(), thrown);
                    break;
                case TRANSACTION_ROLLBACK_FAILED:
                    finalizeException = new TransactionRollbackFailedException(thrown.getMessage(),
                            error.getTransactionRollbackFailed().getReason(), thrown);
                    break;
            }
        }
        if (finalizeException == null) {
            //Either there was no thrown exception, or it didn't have an RpcGitError, or the RpcGitError
            //it had wasn't related to transaction handling. Fall back on using our local state to guess
            //whether we would have tried to commit or roll back
            finalizeException = chooseFinalizeError(thrown);
        }

        if (callback != null && thrown != null) {
            //Notify the callback that this replica transaction has hit an unrecoverable error. This will allow
            //the LocalTransactionLeader to remove the replica from the write lock. The usual mechanism for this
            //is through the prepareStage completing exceptionally, but there is an edge case where the replica
            //has prepared but has not granted its write token yet (or has retracted it since). In this case,
            //the prepareStage has completed, but the write lock is still waiting for the replica's write token
            //(which in case of a unrecoverable error will never be granted).
            callback.onReplicaError(replica);
        }
        finalizeStage.completeExceptionally(finalizeException);
    }

    private class RpcPackWriter implements SatiableConsumer<RpcStreamChunk> {

        private final AtomicBoolean aborted;

        RpcPackWriter(Repository repository) {
            aborted = new AtomicBoolean();

            requestStream.onNext(RpcTransactLeaderFragment.newBuilder()
                    .setLogicalTimestamp(logicalClock.tick())
                    .setStagePackRequest(RpcStagePackRequest.newBuilder()
                            .setTransactionId(transactionId)
                            .setRepository(repository.getId()))
                    .build());
            backoffObserver.activate(); //Apply backpressure to pack chunks to prevent excessive buffering
        }

        public void abort() {
            if (aborted.compareAndSet(false, true)) {
                log.warn("[{}] Aborted pack replication to the follower", this);
                backoffObserver.deactivate(); //Release backpressure, since no further pack chunks will be sent
            }
        }

        @Override
        public boolean accept(@Nonnull RpcStreamChunk chunk) {
            if (completed.get() || aborted.get()) {
                return true;
            }
            requestStream.onNext(RpcTransactLeaderFragment.newBuilder()
                    .setPackChunk(chunk)
                    .setLogicalTimestamp(logicalClock.get())
                    .build());
            return isDone();
        }

        @Nonnull
        @Override
        public String toString() {
            return ReplicaTransactionResponseObserver.this.toString();
        }
    }
}
