package com.atlassian.bitbucket.mesh.git;

import com.atlassian.bitbucket.mesh.coordination.LogicalClock;
import com.atlassian.bitbucket.mesh.coordination.WriteToken;
import com.atlassian.bitbucket.mesh.coordination.WriteTokenHandler;
import com.atlassian.bitbucket.mesh.coordination.WriteTokenRegistry;
import com.atlassian.bitbucket.mesh.event.EventPublisher;
import com.atlassian.bitbucket.mesh.execution.ExecutionManager;
import com.atlassian.bitbucket.mesh.git.exception.InvalidArgumentException;
import com.atlassian.bitbucket.mesh.network.MeshManager;
import com.atlassian.bitbucket.mesh.repository.*;
import com.atlassian.bitbucket.mesh.rpc.v1.*;
import io.grpc.stub.StreamObserver;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.annotation.Nonnull;
import java.util.List;
import java.util.concurrent.atomic.AtomicBoolean;

import static com.atlassian.bitbucket.mesh.repository.ReplicaState.*;
import static com.google.common.base.Preconditions.checkState;

/**
 * {@link StreamObserver} on the follower side of a {@code TransactionService/VerifyConsistency} RPC. This RPC is used
 * to check and synchronize the replica states when it's suspected that they might be internally inconsistent, e.g.
 * after one or more replica nodes have recovered from a backup, or failover.
 * <p>
 * The verify-consistency 'protocol' for a repository is as follows:
 * <ul>
 *     <li>Leader starts a {@code TransactionService/VerifyConsistency} for each remote replica.</li>
 *     <li>Leader sends a token request.</li>
 *     <li>Follower acquires write token and sends a {@link RpcWriteTokenGranted} message to the client, which
 *     includes the follower's repository hashes and version.</li>
 *     <li>When the leader has received all hashes and versions from all replicas, it determines the 'blessed'
 *     repository state and sends a summary to all replicas detailing which of the replicas are inconsistent.</li>
 *     <li>Follower marks itself inconsistent if its listed as inconsistent in the summary.</li>
 *     <li>Follower releases the write token for the repository.</li>
 * </ul>
 *
 * @since 1.2
 */
class VerifyConsistencyFollower implements StreamObserver<RpcVerifyConsistencyLeaderFragment> {

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

    private final AtomicBoolean completed;
    private final EventPublisher eventPublisher;
    private final LogicalClock logicalClock;
    private final MeshManager meshManager;
    private final RepositoryAttributeManager repositoryAttributeManager;
    private final RepositoryHashService repositoryHashService;
    private final RepositoryManager repositoryManager;
    private final StreamObserver<RpcVerifyConsistencyFollowerFragment> responseObserver;
    private final TokenHandler writeTokenHandler;
    private final WriteTokenRegistry writeTokenRegistry;

    private volatile long logicalTimestamp;
    private volatile Repository repository;
    private volatile WriteToken.Promise writeTokenPromise;
    private volatile RpcWriteTokenRequest writeTokenRequest;

    VerifyConsistencyFollower(StreamObserver<RpcVerifyConsistencyFollowerFragment> responseObserver,
                              EventPublisher eventPublisher, ExecutionManager executionManager,
                              LogicalClock logicalClock, MeshManager meshManager,
                              RepositoryAttributeManager repositoryAttributeManager,
                              RepositoryHashService repositoryHashService,
                              RepositoryManager repositoryManager, WriteTokenRegistry writeTokenRegistry) {
        this.eventPublisher = eventPublisher;
        this.logicalClock = logicalClock;
        this.meshManager = meshManager;
        this.repositoryAttributeManager = repositoryAttributeManager;
        this.repositoryHashService = repositoryHashService;
        this.repositoryManager = repositoryManager;
        this.responseObserver = responseObserver;
        this.writeTokenRegistry = writeTokenRegistry;

        completed = new AtomicBoolean();
        writeTokenHandler = new TokenHandler();

        //Ensure the token is released and the promise canceled if the context is closed unexpectedly
        executionManager.getContext().ifPresent(context -> maybeReleaseWriteToken());
    }

    @Override
    public void onCompleted() {
    }

    @Override
    public void onError(Throwable t) {
        //If the error was sent by us, completed is already set to true. If the error was sent by the leader, we
        //still need to complete the RPC from our end
        maybeComplete();

        //Release the write token if held
        maybeReleaseWriteToken();
    }

    @Override
    public void onNext(RpcVerifyConsistencyLeaderFragment fragment) {
        logicalTimestamp = logicalClock.onReceived(fragment.getLogicalTimestamp());
        switch (fragment.getFragmentOneofCase()) {
            case TOKEN_REQUEST:
                onTokenRequest(fragment.getTokenRequest(), fragment.getLogicalTimestamp());
                break;
            case TOKEN_YIELDED:
                //Leader has yielded the write token as requested by us, propagate to the registry
                writeTokenHandler.onYielded();
                break;
            case CONSISTENCY_SUMMARY:
                onSummary(fragment.getConsistencySummary());
                break;
            default:
                throw new InvalidArgumentException("Unsupported fragment type. Only TOKEN_REQUEST, TOKEN_YIELDED," +
                        " and CONSISTENCY_SUMMARY are supported");
        }
    }

    private void maybeComplete() {
        if (completed.compareAndSet(false, true)) {
            responseObserver.onCompleted();
        }
    }

    private void maybeReleaseWriteToken() {
        WriteToken.Promise writeTokenPromise = this.writeTokenPromise;
        if (writeTokenPromise != null) {
            writeTokenPromise.close();
            this.writeTokenPromise = null;
        }
    }

    private void onSummary(RpcRepositoryConsistencySummary summary) {
        try {
            maybeComplete();
            List<String> consistentReplicas = summary.getConsistentReplicasList();
            List<String> inconsistentReplicas = summary.getInconsistentReplicasList();
            List<String> unavailableReplicas = summary.getUnavailableReplicasList();

            String nodeId = meshManager.getLocalNode().getId();
            ReplicaState state = inconsistentReplicas.contains(nodeId) ? INCONSISTENT :
                    consistentReplicas.contains(nodeId) ? CONSISTENT : UNKNOWN;
            if (state == UNKNOWN) {
                log.warn("[{}] This node ({}) was not marked consistent or inconsistent! " +
                        "Not updating replica state..", repository, nodeId);
            } else {
                //Force-update the local replica state while holding the write token
                eventPublisher.publish(new ReplicaStateObservedEvent(writeTokenRequest.getRepository(),
                        new ReplicaStateObservation(state, logicalTimestamp, summary.getVersion(), true)));
            }

            if (inconsistentReplicas.isEmpty() && unavailableReplicas.isEmpty()) {
                log.debug("[{}] Consistency verification completed. All replicas are consistent at " +
                                "(content-hash={}, version={})",
                        repository, summary.getContentHash(), summary.getVersion());
            } else {
                log.info("[{}] Consistency verification is complete. This node is {} with (content-hash: {}, " +
                                "version:{}); consistent-replicas: {}, inconsistent-replicas: {}, " +
                                "unavailable-replicas: {}",
                        repository, state, summary.getContentHash(), summary.getVersion(),
                        consistentReplicas, inconsistentReplicas, unavailableReplicas);
            }
        } finally {
            maybeReleaseWriteToken();
        }
    }

    private void onTokenRequest(RpcWriteTokenRequest request, long tokenTimestamp) {
        checkState(writeTokenRequest == null, "Token has already been requested!");
        writeTokenRequest = request;

        repository = repositoryManager.findById(request.getRepository(), true);
        if (repository.exists()) {
            //Clear the content hash to ensure it is calculated from scratch
            repositoryHashService.resetContentHash(repository);
        } else {
            //If the repository is missing on disk, make sure any trace of it in the database is gone too
            repositoryAttributeManager.delete(request.getRepository());
        }
        //Acquire the token; the writeTokenHandler will send the token granted message when the token is granted
        writeTokenPromise = writeTokenRegistry.acquireToken(
                repository, request.getTransactionId(), tokenTimestamp, writeTokenHandler);
    }

    private void sendFragmentToLeader(RpcVerifyConsistencyFollowerFragment.Builder fragment) {
        if (!completed.get()) {
            responseObserver.onNext(fragment.setLogicalTimestamp(logicalClock.tick()).build());
        }
    }

    /**
     * This handler serves a bridge between the local WriteTokenRegistry and the verify consistency operation that
     * requested the write token.
     * <p>
     * A note on concurrency: the onGranted and yieldToken methods are called from the local WriteTokenRegistry,
     * which internally guarantees thread safety. In other words, TokenHandler does not need to perform internal
     * synchronization to prevent out-of-order processing.
     */
    private class TokenHandler implements WriteTokenHandler {

        private volatile WriteToken token;

        @Override
        public void onGranted(@Nonnull WriteToken token) {
            this.token = token;

            //Now that we hold the lock, send this node's replica version details to the leader
            boolean exists = repository.exists();
            //If the repository does not exist locally, don't even try to calculate the content hash
            String contentHash = exists ? repositoryHashService.getContentHash(repository) : RefHasher.NULL_HASH;

            sendFragmentToLeader(RpcVerifyConsistencyFollowerFragment.newBuilder()
                    .setTokenGranted(RpcWriteTokenGranted.newBuilder()
                            .setReplicaVersion(RpcReplicaVersion.newBuilder()
                                    .setContentHash(contentHash)
                                    .setVersion(repositoryAttributeManager.getVersion(repository.getId())))));
        }

        @Override
        public void yieldToken() {
            //Called by the local WriteTokenRegistry when the token should be yielded.
            //Forward the request to the verification leader.
            sendFragmentToLeader(RpcVerifyConsistencyFollowerFragment.newBuilder()
                    .setYieldTokenRequest(RpcYieldWriteTokenRequest.getDefaultInstance()));
        }

        void onYielded() {
            //Called when the verification leader has yielded the token.
            WriteToken token = this.token;
            if (token != null) {
                //Release the token back to the registry.
                token.yieldToken();
                this.token = null;
            }
        }
    }
}
