package com.atlassian.bitbucket.mesh.transaction;

import com.atlassian.bitbucket.mesh.network.MeshManager;
import com.atlassian.bitbucket.mesh.network.MeshNode;
import com.atlassian.bitbucket.mesh.network.PartitionReplicaMigration;
import com.atlassian.bitbucket.mesh.repository.ManagedRepository;
import com.atlassian.bitbucket.mesh.repository.Repository;
import com.google.common.collect.ImmutableSet;

import javax.annotation.Nonnull;
import java.util.*;
import java.util.stream.Stream;

import static com.atlassian.bitbucket.mesh.util.MoreCollectors.toImmutableSet;

/**
 * Helper class used by {@link LocalTransactionLeader} to determine which combinations of {@link ReplicaTransaction
 * replica transactions} need to reach a write quorum before a coordinated write is allowed to commit. The logic
 * takes ongoing migrations into account, and has been implemented in a helper class to allow easier unit testing
 * of this logic.
 */
class WriteQuorumHelper {

    private final MeshManager meshManager;

    WriteQuorumHelper(MeshManager meshManager) {
        this.meshManager = meshManager;
    }

    @Nonnull
    Set<Set<MeshNode>> getQuorumSets(@Nonnull Repository repository,
                                     @Nonnull List<ReplicaTransaction> replicaTransactions) {
        Set<MeshNode> replicas = Stream.concat(Stream.of(meshManager.getLocalNode()), replicaTransactions.stream()
                .map(ReplicaTransaction::getNode))
                .collect(toImmutableSet());
        Collection<PartitionReplicaMigration> migrations = repository instanceof ManagedRepository
                ? meshManager.getPartitionMigrations(((ManagedRepository) repository).getPartition())
                : Collections.emptyList(); //only managed repositories can have partition replica migrations
        if (migrations.isEmpty()) {
            //There are no partition migrations, return a single quorum set of the replicas + the local node
            return ImmutableSet.of(replicas);
        }

        List<Set<MeshNode>> quorumSets = new LinkedList<>();
        quorumSets.add(replicas);

        for (PartitionReplicaMigration migration : migrations) {
            //Every migration has two possible outcomes:
            // - if it completes, the migration's target node holds a partition replica and the source node does not
            // - if it fails or is canceled, the source node holds a partition replica and the target node does not
            //This loop generates a quorum set for each outcome, by taking the quorum sets we've determined so far and
            //applying each outcome to it by _removing_ the node from the quorum that would not hold a replica in the
            //possible outcome.
            ListIterator<Set<MeshNode>> it = quorumSets.listIterator();
            while (it.hasNext()) {
                Set<MeshNode> quorumSet = it.next();
                MeshNode source = migration.getSourceNode();
                MeshNode target = migration.getTargetNode();
                if (!replicas.contains(source) || !replicas.contains(target)) {
                    //The migration references a node that we don't consider to hold a replica. This can happen
                    //if a topology update is received from the control plane _after_ the replica transactions have
                    //been created but _before_ the quorum sets are determined. In that case we should ignore the
                    //migration.
                    continue;
                }
                if (quorumSet.contains(source) || quorumSet.contains(target)) {
                    //Replica the current scenario with two scenarios: one where the migration completes and one
                    //where the migration rolls back
                    it.remove();
                    it.add(copyOfExcluding(quorumSet, source));
                    it.add(copyOfExcluding(quorumSet, target));
                }
            }
        }
        //Deduplicate possible outcomes
        return ImmutableSet.copyOf(quorumSets);
    }

    private static Set<MeshNode> copyOfExcluding(Set<MeshNode> sourceSet, MeshNode value) {
        Set<MeshNode> result = new HashSet<>(sourceSet);
        result.remove(value);
        return result;
    }
}
