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.BitbucketServerRepository;
import com.atlassian.bitbucket.mesh.repository.ManagedRepository;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableSet;
import com.google.common.collect.Sets;
import org.junit.Before;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.mockito.Mock;
import org.mockito.junit.MockitoJUnitRunner;

import java.util.*;

import static java.util.stream.Collectors.toSet;
import static org.hamcrest.MatcherAssert.assertThat;
import static org.hamcrest.Matchers.containsInAnyOrder;
import static org.junit.Assert.assertNotNull;
import static org.mockito.ArgumentMatchers.anyString;
import static org.mockito.Mockito.*;

@RunWith(MockitoJUnitRunner.class)
public class WriteQuorumHelperTest {

    private static final String PARTITION = "0008";

    private WriteQuorumHelper helper;
    @Mock
    private MeshManager meshManager;
    private Set<PartitionReplicaMigration> mockedPartitionMigrations;
    private List<MeshNode> nodes;
    @Mock
    private ManagedRepository repository;

    @Before
    public void setUp() {
        mockedPartitionMigrations = new HashSet<>();
        when(meshManager.getPartitionMigrations(anyString())).thenAnswer(invocation -> {
            String partition = invocation.getArgument(0);
            return PARTITION.equals(partition) ? mockedPartitionMigrations : ImmutableSet.of();
        });
        when(repository.getPartition()).thenReturn(PARTITION);
        nodes = new ArrayList<>();
        for (int i = 0; i < 6; ++i) {
            MeshNode node = mock(MeshNode.class);
            lenient().when(node.toString()).thenReturn("node-" + i); // helpful when debugging
            nodes.add(node);
        }
        //node0 is the local node
        when(meshManager.getLocalNode()).thenReturn(nodes.get(0));
        helper = new WriteQuorumHelper(meshManager);
    }

    @Test
    public void testMigrationStartedAfterReplicationStarted() {
        //Mock a scenario where we're replicating writes to 3 replicas (nodes 1-3), and we have partition
        //migration 0->3 and a migration 1->4 that is started _after_ the replicated write was started.
        //This can happen if Mesh receives a topology update after a write replication was started, but _before_
        //the quorum sets are calculated.
        //In this case we want the migration to an unknown replica to be ignored.
        List<ReplicaTransaction> replicaTransactions =
                mockReplicaTransactions(nodes.get(1), nodes.get(2), nodes.get(3));
        mockPartitionMigration(nodes.get(0), nodes.get(3));
        mockPartitionMigration(nodes.get(1), nodes.get(4));

        Set<Set<MeshNode>> result = helper.getQuorumSets(repository, replicaTransactions);
        assertNotNull(result);
        assertThat(result, containsInAnyOrder(
                nodes(0, 1, 2), // 0->3 migration completes
                nodes(1, 2, 3))); // 0->3 migration canceled
    }

    @Test
    public void testNoMigrations() {
        List<ReplicaTransaction> replicaTransactions =
                mockReplicaTransactions(nodes.get(1), nodes.get(2));
        Set<Set<MeshNode>> result = helper.getQuorumSets(repository, replicaTransactions);
        assertNotNull(result);
        assertThat(result, containsInAnyOrder(nodes(0, 1, 2)));
    }

    @Test
    public void testNonManagedRepository() {
        //This test verifies a scenario that should never occur to ensure it handles the unexpected call gracefully;
        //only managed repositories should trigger coordinated writes, so getQuorumSets should never be called with
        //a non-empty list of replica transactions, and a non-managed repository.
        List<ReplicaTransaction> replicaTransactions = mockReplicaTransactions(nodes.get(1), nodes.get(2));

        BitbucketServerRepository classicRepository = mock(BitbucketServerRepository.class);
        Set<Set<MeshNode>> result = helper.getQuorumSets(classicRepository, replicaTransactions);
        assertNotNull(result);
        assertThat(result, containsInAnyOrder(nodes(0, 1, 2)));
    }

    @Test
    public void testSingleMigration() {
        //Mock a scenario where we're replicating writes to 3 replicas, and we have partition migration 1 -> 3
        MeshNode source = nodes.get(1);
        MeshNode target = nodes.get(3);
        List<ReplicaTransaction> replicaTransactions = mockReplicaTransactions(source, nodes.get(2), target);
        mockPartitionMigration(source, target);

        Set<Set<MeshNode>> result = helper.getQuorumSets(repository, replicaTransactions);
        assertNotNull(result);
        assertThat(result, containsInAnyOrder(
                nodes(0, 1, 2), // migration canceled
                nodes(0, 2, 3))); // migration completed
    }

    @Test
    public void testThreeMigrations() {
        //Mock a scenario where we're replicating writes to 5 replicas, and we have partition migrations
        //from 0 -> 3, 1 -> 4 and 2 -> 5
        MeshNode localNode = nodes.get(0);
        MeshNode node1 = nodes.get(1);
        MeshNode node2 = nodes.get(2);
        MeshNode node3 = nodes.get(3);
        MeshNode node4 = nodes.get(4);
        MeshNode node5 = nodes.get(5);
        List<ReplicaTransaction> replicaTransactions = mockReplicaTransactions(node1, node2, node3, node4, node5);
        mockPartitionMigration(localNode, node3);
        mockPartitionMigration(node1, node4);
        mockPartitionMigration(node2, node5);

        Set<Set<MeshNode>> result = helper.getQuorumSets(repository, replicaTransactions);
        assertNotNull(result);
        assertThat(result, containsInAnyOrder(
                nodes(3, 4, 5), //0->3, 1->4 and 2->5 complete
                nodes(2, 3, 4), //0->3, 1->4 complete and 2->5 canceled
                nodes(1, 3, 5), //0->3 and 2->5 complete, 1->4 canceled
                nodes(1, 2, 3), //0->3 complete, 1->4 and 2->5 canceled
                nodes(0, 4, 5), //0->3 canceled, 1->4 and 2->5 complete
                nodes(0, 2, 4), //1->4 complete, 0->3 and 2->5 canceled
                nodes(0, 1, 5), //2->5 complete, 0->3, 1->4 canceled
                nodes(0, 1, 2))); //0->3, 1->4 and 2->5 canceled
    }

    @Test
    public void testTwoMigrations() {
        //Mock a scenario where we're replicating writes to 4 replicas, and we have partition migrations
        //from 1 -> 3 and 2 -> 4
        MeshNode node1 = nodes.get(1);
        MeshNode node2 = nodes.get(2);
        MeshNode node3 = nodes.get(3);
        MeshNode node4 = nodes.get(4);
        List<ReplicaTransaction> replicaTransactions = mockReplicaTransactions(node1, node2, node3, node4);
        mockPartitionMigration(node1, node3);
        mockPartitionMigration(node2, node4);

        Set<Set<MeshNode>> result = helper.getQuorumSets(repository, replicaTransactions);
        assertNotNull(result);
        assertThat(result, containsInAnyOrder(
                nodes(0, 3, 4), //1-> 3 and 2 -> 4 complete
                nodes(0, 2, 3), //1-> 3 completes, 2 -> 4 canceled
                nodes(0, 1, 4), //1-> 3 canceled, 2 -> 4 completes
                nodes(0, 1, 2))); //1-> 3 and 2 -> 4 canceled
    }

    private static Set<ReplicaTransaction> copyExcluding(Collection<ReplicaTransaction> original,
                                                         MeshNode... values) {
        Set<MeshNode> excludes = Sets.newHashSet(values);
        return original.stream()
                .filter(replicaTransaction -> !excludes.contains(replicaTransaction.getNode()))
                .collect(toSet());
    }

    private static List<ReplicaTransaction> mockReplicaTransactions(MeshNode... nodes) {
        ImmutableList.Builder<ReplicaTransaction> builder = ImmutableList.builder();
        for (MeshNode node : nodes) {
            ReplicaTransaction transaction = mock(ReplicaTransaction.class);
            when(transaction.getNode()).thenReturn(node);
            String transactionName = "replica-transaction-" + node;
            lenient().when(transaction.toString()).thenReturn(transactionName); //easier debugging
            builder.add(transaction);
        }
        return builder.build();
    }

    private void mockPartitionMigration(MeshNode source, MeshNode target) {
        mockedPartitionMigrations.add(new PartitionReplicaMigration.Builder()
                .partition(PARTITION)
                .sourceNode(source)
                .targetNode(target)
                .build());
    }

    private Set<MeshNode> nodes(int... indices) {
        ImmutableSet.Builder<MeshNode> builder = ImmutableSet.builder();
        for (int index : indices) {
            builder.add(nodes.get(index));
        }
        return builder.build();
    }
}