package com.atlassian.bitbucket.mesh.transaction;

import com.atlassian.bitbucket.mesh.AbstractStatusException;
import com.atlassian.bitbucket.mesh.client.MeshClient;
import com.atlassian.bitbucket.mesh.client.ReplicateParameters;
import com.atlassian.bitbucket.mesh.coordination.LogicalClock;
import com.atlassian.bitbucket.mesh.coordination.WriteLock;
import com.atlassian.bitbucket.mesh.coordination.WriteTokenRegistry;
import com.atlassian.bitbucket.mesh.event.EventPublisher;
import com.atlassian.bitbucket.mesh.git.exception.AbortedException;
import com.atlassian.bitbucket.mesh.git.exception.FailedPreconditionException;
import com.atlassian.bitbucket.mesh.git.exception.InternalFailureException;
import com.atlassian.bitbucket.mesh.gprc.GrpcSupport;
import com.atlassian.bitbucket.mesh.management.SidebandManager;
import com.atlassian.bitbucket.mesh.network.MeshManager;
import com.atlassian.bitbucket.mesh.network.MeshNode;
import com.atlassian.bitbucket.mesh.network.Partition;
import com.atlassian.bitbucket.mesh.network.PartitionReplica;
import com.atlassian.bitbucket.mesh.repository.*;
import com.atlassian.bitbucket.mesh.rpc.v1.RpcReplicaState;
import com.atlassian.bitbucket.mesh.rpc.v1.RpcRepositoryConsistencySummary;
import com.atlassian.bitbucket.mesh.rpc.v1.RpcRepositoryReplicaObservation;
import com.atlassian.bitbucket.mesh.rpc.v1.RpcSidebandFollowerFragment;
import com.atlassian.bitbucket.mesh.util.ShaUtils;
import com.google.common.base.Throwables;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.Iterables;
import io.grpc.Status;
import io.grpc.StatusRuntimeException;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.annotation.Nonnull;
import javax.annotation.Nullable;
import java.time.Duration;
import java.util.*;
import java.util.concurrent.*;
import java.util.stream.Stream;

import static com.google.common.base.Preconditions.checkArgument;
import static com.google.common.collect.ImmutableList.toImmutableList;
import static com.google.common.collect.ImmutableSet.toImmutableSet;
import static java.util.Objects.requireNonNull;

/**
 * Verifies the consistency of all replicas of the provided repository. This is use when one or more sibling nodes
 * has been restored from backup, or the system has failed over to standby nodes that may not all be synced to the
 * exact same state.
 * <ul>
 *     <li>Requesting the write token from each of the replicas, including the local replica.</li>
 *     <li>Each replica returns its write token, and provides its local version and hashes.</li>
 *     <li>The leader (this node) compares all versions and hashes and selects the 'blessed' version and hashes.</li>
 *     <li>The leader then determines which replicas, if any, are inconsistent and returns a summary to all peers.</li>
 *     <li>Each replica updates its local replica state, schedules a repair operation if inconsistent, and releases its
 *     write token.</li>
 *     <li>The leader updates the control plane with each replica's state.</li>
 * </ul>
 *
 * @since 1.2
 */
class RepositoryConsistencyVerifier {

    private static final Logger log = LoggerFactory.getLogger(RepositoryConsistencyVerifier.class);
    private static final long DEFAULT_TIMEOUT_MS = Duration.ofMinutes(1L).toMillis();

    private final EventPublisher eventPublisher;
    private final LogicalClock logicalClock;
    private final MeshClient meshClient;
    private final MeshManager meshManager;
    private final ConcurrentMap<MeshNode, ReplicaVersion> replicaVersions;
    private final RepositoryAttributeManager repositoryAttributeManager;
    private final RepositoryHashService repositoryHashService;
    private final String repositoryId;
    private final RepositoryManager repositoryManager;
    private final SidebandManager sidebandManager;
    private final String transactionId;
    private final WriteTokenRegistry writeTokenRegistry;

    private Set<MeshNode> replicas;
    private List<ReplicaConsistencyVerifier> replicaVerifiers;
    private volatile Repository repository;
    private volatile RepositoryConsistencySummary summary;
    private volatile WriteLock writeLock;

    RepositoryConsistencyVerifier(String repositoryId, String transactionId, LogicalClock logicalClock,
                                  EventPublisher eventPublisher, MeshClient meshClient, MeshManager meshManager,
                                  RepositoryAttributeManager repositoryAttributeManager,
                                  RepositoryHashService repositoryHashService, RepositoryManager repositoryManager,
                                  SidebandManager sidebandManager, WriteTokenRegistry writeTokenRegistry) {
        this.eventPublisher = eventPublisher;
        this.logicalClock = logicalClock;
        this.meshClient = meshClient;
        this.meshManager = meshManager;
        this.repositoryAttributeManager = repositoryAttributeManager;
        this.repositoryHashService = repositoryHashService;
        this.repositoryId = repositoryId;
        this.repositoryManager = repositoryManager;
        this.sidebandManager = sidebandManager;
        this.transactionId = transactionId;
        this.writeTokenRegistry = writeTokenRegistry;

        replicaVersions = new ConcurrentHashMap<>();
    }

    @Nullable
    RepositoryConsistencySummary call() {
        ManagedRepository.Id id = ManagedRepository.Id.parseId(repositoryId).orElse(null);
        if (id == null) {
            log.debug("[{}] Skipping replica consistency checks for non-managed repository", repositoryId);
            return null;
        }
        Partition partition = meshManager.getPartition(id.getPartition());
        if (partition == null) {
            throw new FailedPreconditionException(String.format(
                    "Cannot check replica consistency for %s; cannot determine peers for unknown partition %s",
                    repositoryId, id.getPartition()));
        }
        replicas = partition.getReplicas().stream()
                .map(PartitionReplica::getNode)
                .filter(MeshNode::isRemote)
                .collect(toImmutableSet());

        log.debug("[{}] Starting replica consistency check (tx = {})", repositoryId, transactionId);
        long lockTimestamp = logicalClock.tick();
        repository = repositoryManager.findById(repositoryId, true);
        try {
            //Start the lock
            log.debug("[{}] Starting write lock for transaction {}", repository, transactionId);
            writeLock = writeTokenRegistry.startLock(repository, transactionId, replicas);
            startReplicaVerifiers(replicas, lockTimestamp);

            //When all replicas have acquired the write lock and sent back their local replica hashes + version,
            //compare them all and determine whether any replicas (including the local replica) are inconsistent.
            long timeoutMs = GrpcSupport.getTimeout().toMillis();
            if (timeoutMs == 0L) {
                //No RPC deadline has been set, or not invoked from an RPC
                timeoutMs = DEFAULT_TIMEOUT_MS;
            }
            writeLock.asCompletionStage()
                    .thenApply(this::onLockAcquired)
                    .toCompletableFuture()
                    .get(timeoutMs, TimeUnit.MILLISECONDS);

            return summary;
        } catch (InterruptedException interruptedException) {
            log.warn("[{}] Interrupted while verifying replica consistency", repository);
            Thread.currentThread().interrupt();
            throw new AbortedException("Aborted while waiting for write lock to be acquired", interruptedException);
        } catch (ExecutionException executionException) {
            Throwable cause = executionException.getCause();
            //A RepositoryNotFoundException is expected and can be propagated unchanged. It means that the repository
            //was missing on all replicas.
            Throwables.throwIfInstanceOf(cause, RepositoryNotFoundException.class);
            //Unexpected exception
            log.warn("[{}] Failed to verify replica consistency", repository, cause);
            Throwables.throwIfInstanceOf(cause, StatusRuntimeException.class);
            Throwables.throwIfInstanceOf(cause, AbstractStatusException.class);
            throw new InternalFailureException("Failed to verify consistency of " + repositoryId, cause);
        } catch (TimeoutException e) {
            log.warn("[{}] Replica consistency verification timed out", repository);
            throw Status.DEADLINE_EXCEEDED
                    .withDescription("Timed out while waiting for write lock to be acquired")
                    .asRuntimeException();
        } finally {
            wrapUp();
        }
    }

    private RepositoryConsistencySummary createSummary(ReplicaVersion blessedVersion) {
        RepositoryConsistencySummary.Builder builder = new RepositoryConsistencySummary.Builder(blessedVersion);

        Stream.concat(Stream.of(meshManager.getLocalNode()), replicas.stream()).forEach(node -> {
            ReplicaVersion nodeVersion = replicaVersions.get(node);
            if (nodeVersion == null) {
                //The node was offline or failed to report its version
                builder.unavailableReplicas(node);
            } else if (nodeVersion.equals(blessedVersion)) {
                builder.consistentReplicas(node);
            } else {
                builder.inconsistentReplicas(node);
            }
        });
        return builder.build();
    }

    private void notifyControlPlaneOfInconsistentReplicas(RepositoryConsistencySummary summary) {
        if (!summary.getInconsistentReplicas().isEmpty()) {
            sidebandManager.send(RpcSidebandFollowerFragment.newBuilder()
                    .setRepositoryReplicaObservation(RpcRepositoryReplicaObservation.newBuilder()
                            .setRepository(repositoryId)
                            .setState(RpcReplicaState.REPLICA_STATE_INCONSISTENT)
                            .setVersion(summary.getConsensusVersion().getVersion())
                            .addAllNodes(Iterables.transform(summary.getInconsistentReplicas(), MeshNode::getId)))
                    .build());
        }
    }

    private Void onLockAcquired(Void result) {
        //All replicas reported their state; we can determine which node(s) have the most up-to-date version
        ReplicaVersion localVersion;
        if (repository.exists()) {
            //Clear the content hash to force a fresh calculation
            repositoryHashService.resetContentHash(repository);
            localVersion = new ReplicaVersion(
                    repositoryHashService.getContentHash(repository),
                    repositoryAttributeManager.getVersion(repositoryId));
        } else {
            //If the repository does not exist on disk, clear out any trace of it from the database too
            repositoryAttributeManager.delete(repositoryId);
            localVersion = new ReplicaVersion(RefHasher.NULL_HASH, 1L);
        }

        replicaVersions.put(meshManager.getLocalNode(), localVersion);
        ReplicaVersion consensusVersion = selectConsensusVersion(replicaVersions.values());

        //Build the summary and send it to each of the replicas
        summary = createSummary(consensusVersion);
        long timestamp = logicalClock.tick();
        notifyControlPlaneOfInconsistentReplicas(summary);

        RpcRepositoryConsistencySummary rpcSummary = summary.toRpc();
        replicaVerifiers.forEach(verifier -> verifier.summarize(rpcSummary, timestamp));
        replicaVerifiers = ImmutableList.of();

        //Update the local state registry
        ReplicaState state = localVersion.equals(consensusVersion) ? ReplicaState.CONSISTENT : ReplicaState.INCONSISTENT;
        eventPublisher.publish(new ReplicaStateObservedEvent(repositoryId,
                new ReplicaStateObservation(state, timestamp, consensusVersion.getVersion(), true)));

        if (summary.getInconsistentReplicas().isEmpty() && summary.getUnavailableReplicas().isEmpty()) {
            if (consensusVersion.getVersion() < 2L && ShaUtils.isNull(consensusVersion.getContentHash())) {
                //Repository creation leads to a version of 2 (default value of 1 which is incremented when
                //repository creation is successfully replicated). Therefore any version smaller than 2
                log.info("[{}] Consistency verification completed. The repository is missing on all replicas!",
                        repository);
                throw new RepositoryNotFoundException(repositoryId, "The repository is missing on all replicas");
            }
            log.debug("[{}] Consistency verification completed. All replicas are consistent at {}",
                    repository, consensusVersion);
        } else {
            log.info("[{}] Consistency verification is complete: consensus-version: '{}', " +
                            "consistent-replicas: {}, inconsistent-replicas: {}, unavailable-replicas: {}",
                    repository, consensusVersion, summary.getConsistentReplicas(),
                    summary.getInconsistentReplicas(), summary.getUnavailableReplicas());
        }
        return null;
    }

    private void startReplicaVerifiers(Set<MeshNode> replicas, long lockTimestamp) {
        ReplicateParameters parameters = new ReplicateParameters.Builder(meshManager.getLocalNode(), transactionId)
                .callback(new TokenCallback())
                .build();
        //Create a verifier for each remote replica
        replicaVerifiers = replicas.stream()
                .map(node -> meshClient.verifyConsistency(node, parameters))
                .collect(toImmutableList());

        //Start the verification process
        replicaVerifiers.forEach(verifier -> verifier.requestTokenAndReplicaVersion(repositoryId, lockTimestamp)
                .handle((result, throwable) -> {
                    if (throwable != null) {
                        Status status = Status.fromThrowable(throwable);
                        MeshNode node = verifier.getReplica();
                        if (status.getCode() == Status.UNAVAILABLE.getCode()) {
                            log.info("[{}] Replica {} is unavailable for consistency check", repositoryId, node);
                        } else {
                            log.warn("[{}] Replica {} failed consistency verification!", repositoryId, node, throwable);
                        }
                        //Remove the replica from the write lock; we're never going to receive its write token
                        log.info("[{}] Removing replica {} from the write lock", repository, node);
                        writeLock.remove(node);
                    }
                    return result;
                }));
    }

    /**
     * Returns the {@link ReplicaVersion} with the highest {@code version}. If there are different
     * {@link ReplicaVersion}s at that same version, but with different content hashes, we select
     * the content hash that occurs most frequently.
     */
    private static ReplicaVersion selectConsensusVersion(Collection<ReplicaVersion> replicaVersions) {
        checkArgument(!replicaVersions.isEmpty(), "replicaVersions cannot be empty");
        long maxVersion = 0L;
        String bestContentHash = null;
        int bestContentHashCount = 0;
        Map<String, Integer> countByContentHash = new HashMap<>();
        for (ReplicaVersion replicaVersion : replicaVersions) {
            if (replicaVersion.getVersion() > maxVersion || ShaUtils.isNull(bestContentHash)) {
                maxVersion = replicaVersion.getVersion();
                countByContentHash.clear();
                countByContentHash.put(replicaVersion.getContentHash(), 1);
                bestContentHash = replicaVersion.getContentHash();
                bestContentHashCount = 1;
            } else if (replicaVersion.getVersion() == maxVersion && !ShaUtils.isNull(replicaVersion.getContentHash())) {
                String contentHash = replicaVersion.getContentHash();
                int count = countByContentHash.compute(contentHash,
                        (ignored, current) -> current == null ? 1 : current + 1);
                if (contentHash.equals(bestContentHash)) {
                    bestContentHashCount = count;
                } else if (count > bestContentHashCount) {
                    bestContentHash = contentHash;
                    bestContentHashCount = count;
                }
            }
        }

        return new ReplicaVersion(requireNonNull(bestContentHash, "bestContentHash"), maxVersion);
    }

    private void wrapUp() {
        if (!replicaVerifiers.isEmpty()) {
            String message = "Aborting consistency verification";
            replicaVerifiers.forEach(verifier -> verifier.cancel(message));
        }
        if (writeLock != null) {
            writeLock.close();
        }
    }

    private class TokenCallback implements ReplicaTransactionCallback {

        @Override
        public void onReplicaError(@Nonnull MeshNode node) {
            if (summary == null) {
                ReplicaVersion version = replicaVersions.remove(node);
                if (version != null) {
                    log.info("[{}] Dropping reported {} version {} because the replica failed", repository, version, node);
                }
                if (writeLock.remove(node)) {
                    log.info("[{}] Removed failed replica {} from the write lock", repository, node);
                }
            }
        }

        @Override
        public boolean onTokenGranted(@Nonnull MeshNode node, @Nonnull ReplicaVersion replicaVersion) {
            replicaVersions.put(node, replicaVersion);
            log.debug("[{}] Replica {} granted its write token at version '{}'", repository, node, replicaVersion);
            return writeLock.onTokenGranted(node);
        }

        @Override
        public boolean yieldToken(@Nonnull MeshNode node) {
            if (writeLock.yieldToken(node)) {
                replicaVersions.remove(node);
                log.debug("[{}] Yielded write token to replica {} as requested", repository, node);
                return true;
            }
            log.debug("[{}] Not yielding write token to replica {} because the full write lock " +
                    "has already been acquired.", repository, node);
            return false;
        }
    }
}
