package com.atlassian.bitbucket.mesh.transaction;

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.InternalFailureException;
import com.atlassian.bitbucket.mesh.junit.LogListener;
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.RpcRepositoryConsistencySummary;
import com.google.common.base.Throwables;
import com.google.common.collect.ImmutableList;
import io.grpc.Status;
import org.hamcrest.Matchers;
import org.junit.After;
import org.junit.Before;
import org.junit.Rule;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.mockito.ArgumentCaptor;
import org.mockito.Captor;
import org.mockito.Mock;
import org.mockito.junit.MockitoJUnitRunner;

import java.time.Duration;
import java.time.Instant;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.Future;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicLong;
import java.util.concurrent.locks.LockSupport;
import java.util.function.BooleanSupplier;

import static org.hamcrest.MatcherAssert.assertThat;
import static org.hamcrest.Matchers.containsInAnyOrder;
import static org.hamcrest.Matchers.instanceOf;
import static org.junit.Assert.*;
import static org.mockito.ArgumentMatchers.*;
import static org.mockito.Mockito.*;

@RunWith(MockitoJUnitRunner.class)
public class RepositoryConsistencyVerifierTest {

    private static final ReplicaVersion DEFAULT_VERSION = new ReplicaVersion("content-hash", 99L);
    private static final String REPOSITORY_ID = ManagedRepository.formatId(17, "abcde12348deca285eb8", "923");
    private static final long START_TIMESTAMP = 82L;
    private static final Duration TIMEOUT = Duration.ofSeconds(2L);
    private static final String TX_ID = "transaction-id";

    @Rule
    public final LogListener logListener = new LogListener(RepositoryConsistencyVerifier.class);

    private RepositoryConsistencyVerifier consistencyVerifier;
    @Mock
    private EventPublisher eventPublisher;
    @Mock
    private LogicalClock logicalClock;
    @Mock
    private MeshClient meshClient;
    @Mock
    private MeshManager meshManager;
    @Mock
    private MeshNode node1;
    @Mock
    private MeshNode node2;
    @Mock
    private MeshNode node3;
    private ReplicaTransactionCallback replicaTransactionCallback;
    @Mock
    private ManagedRepository repository;
    @Mock
    private RepositoryAttributeManager repositoryAttributeManager;
    @Mock
    private RepositoryHashService repositoryHashService;
    @Mock
    private RepositoryManager repositoryManager;
    @Mock
    private SidebandManager sidebandManager;
    @Captor
    private ArgumentCaptor<RpcRepositoryConsistencySummary> summaryCaptor;
    private Thread verificationThread;
    private Map<MeshNode, ReplicaConsistencyVerifier> verifiersByReplica;
    @Mock
    private WriteLock writeLock;
    private CompletableFuture<Void> writeLockFuture;
    private Map<MeshNode, CompletableFuture<Void>> writeTokenFutureByReplica;
    @Mock
    private WriteTokenRegistry writeTokenRegistry;
    private Set<MeshNode> writeTokensPending;

    @Before
    public void setUp() {
        when(repositoryManager.findById(REPOSITORY_ID, true)).thenReturn(repository);
        when(repository.exists()).thenReturn(true);

        //Mock a partition with 3 replicas, where node2 is the local node
        Partition partition = mock(Partition.class);
        when(meshManager.getPartition("0011")).thenReturn(partition);
        when(partition.getReplicas()).thenReturn(ImmutableList.of(
                new PartitionReplica.Builder("0011-1", node1).build(),
                new PartitionReplica.Builder("0011-2", node2).build(),
                new PartitionReplica.Builder("0011-3", node3).build()));
        when(node1.getId()).thenReturn("1");
        when(node1.isRemote()).thenReturn(true);
        when(node2.getId()).thenReturn("2");
        lenient().when(node2.isLocal()).thenReturn(true);
        when(node3.getId()).thenReturn("3");
        when(node3.isRemote()).thenReturn(true);
        when(meshManager.getLocalNode()).thenReturn(node2);

        //Mock the logical clock
        AtomicLong timestamp = new AtomicLong(START_TIMESTAMP);
        when(logicalClock.tick()).thenAnswer(ignored -> timestamp.incrementAndGet());

        //Set up the ReplicaConsistencyVerifier mocks for each replica.
        verifiersByReplica = new HashMap<>(); //Tracked for verification
        writeTokenFutureByReplica = new HashMap<>(); //Tracked to control write tokens being granted in the tests
        when(meshClient.verifyConsistency(any(), any())).thenAnswer(invocation -> {
            MeshNode replica = invocation.getArgument(0);
            assertFalse(verifiersByReplica.containsKey(replica));
            ReplicateParameters parameters = invocation.getArgument(1);
            replicaTransactionCallback = parameters.getCallback();

            ReplicaConsistencyVerifier result = mock(ReplicaConsistencyVerifier.class);
            lenient().when(result.getReplica()).thenReturn(replica);
            when(result.requestTokenAndReplicaVersion(eq(REPOSITORY_ID), anyLong())).thenAnswer(tokenInvocation -> {
                assertFalse(writeTokenFutureByReplica.containsKey(replica));
                CompletableFuture<Void> tokenFuture = new CompletableFuture<>();
                writeTokenFutureByReplica.put(replica, tokenFuture);
                return tokenFuture;
            });
            verifiersByReplica.put(replica, result);
            return result;
        });

        //Mock the write lock
        writeTokensPending = new HashSet<>();
        writeTokensPending.add(node1);
        writeTokensPending.add(node2);
        writeTokensPending.add(node3);

        when(writeTokenRegistry.startLock(eq(repository), eq(TX_ID), any())).thenReturn(writeLock);
        writeLockFuture = new CompletableFuture<>();
        when(writeLock.asCompletionStage()).thenReturn(writeLockFuture);
        when(writeLock.remove(any())).thenAnswer(invocation -> {
            MeshNode node = invocation.getArgument(0);
            return writeTokensPending.remove(node);
        });
        when(writeLock.yieldToken(any())).thenAnswer(invocation -> {
            MeshNode node = invocation.getArgument(0);
            return writeTokensPending.add(node); // only return true if the token had already been granted
        });

        //Mock local repository state - can be overridden by grantWriteTokens
        when(repositoryAttributeManager.getVersion(REPOSITORY_ID)).thenReturn(DEFAULT_VERSION.getVersion());
        when(repositoryHashService.getContentHash(repository)).thenReturn(DEFAULT_VERSION.getContentHash());

        consistencyVerifier = new RepositoryConsistencyVerifier(REPOSITORY_ID, TX_ID, logicalClock, eventPublisher,
                meshClient, meshManager, repositoryAttributeManager, repositoryHashService, repositoryManager,
                sidebandManager, writeTokenRegistry);
    }

    @After
    public void tearDown() {
        if (verificationThread != null && verificationThread.isAlive()) {
            verificationThread.interrupt();
            verificationThread = null;
        }
    }

    @Test
    public void testAllReplicasConsistent() throws Exception {
        Future<RepositoryConsistencySummary> future = start();

        assertFalse(future.isDone());
        //Grant all write tokens at DEFAULT_VERSION
        grantWriteTokens(DEFAULT_VERSION, node1, node2, node3);

        RepositoryConsistencySummary summary = future.get(TIMEOUT.toMillis(), TimeUnit.MILLISECONDS);
        assertNotNull(summary);

        long tokenRequestTimestamp = START_TIMESTAMP + 1L; // clock is ticked before sending a message
        verify(verifiersByReplica.get(node1)).requestTokenAndReplicaVersion(REPOSITORY_ID, tokenRequestTimestamp);
        verify(verifiersByReplica.get(node3)).requestTokenAndReplicaVersion(REPOSITORY_ID, tokenRequestTimestamp);
        verify(repositoryHashService).resetContentHash(repository);
        long summaryTimestamp = tokenRequestTimestamp + 1L; // clock is ticked before sending a message
        verify(verifiersByReplica.get(node1)).summarize(summaryCaptor.capture(), eq(summaryTimestamp));
        verify(verifiersByReplica.get(node3)).summarize(summaryCaptor.capture(), eq(summaryTimestamp));

        RpcRepositoryConsistencySummary rpcSummary = summaryCaptor.getValue();
        assertEquals(DEFAULT_VERSION.getContentHash(), rpcSummary.getContentHash());
        assertEquals(DEFAULT_VERSION.getVersion(), rpcSummary.getVersion());
        assertEquals(0, rpcSummary.getInconsistentReplicasList().size());
        assertEquals(0, rpcSummary.getUnavailableReplicasList().size());

        verifyReplicaStateObserved(ReplicaState.CONSISTENT, DEFAULT_VERSION.getVersion(), summaryTimestamp);
        verify(writeLock).close();
    }

    @Test
    public void testAllReplicasMissing() {
        Future<RepositoryConsistencySummary> future = start();

        assertFalse(future.isDone());
        //Grant all write tokens with a replica version that represents a missing repository
        grantWriteTokens(new ReplicaVersion(RefHasher.NULL_HASH, 1L), node1, node2, node3);

        ExecutionException exception = assertThrows(ExecutionException.class, () ->
                future.get(TIMEOUT.toMillis(), TimeUnit.MILLISECONDS));
        assertThat(exception.getCause(), instanceOf(RepositoryNotFoundException.class));

        verifyReplicaStateObserved(ReplicaState.CONSISTENT, 1L, START_TIMESTAMP + 2L);
        verify(writeLock).close();
    }

    @Test
    public void testFailedReplicaIsRemovedFromWriteLock() throws Exception {
        Future<RepositoryConsistencySummary> future = start();

        //Grant tokens for 2 of the 3 nodes
        grantWriteTokens(DEFAULT_VERSION, node1, node3);

        //Mock node3 failing (RPC terminates) before the write lock is achieved.
        replicaTransactionCallback.onReplicaError(node3);

        //Grant the last token and verify that the operation terminates
        grantWriteTokens(DEFAULT_VERSION, node2);

        RepositoryConsistencySummary summary = future.get(TIMEOUT.toMillis(), TimeUnit.MILLISECONDS);
        assertEquals(DEFAULT_VERSION, summary.getConsensusVersion());
        assertThat(summary.getConsistentReplicas(), containsInAnyOrder(node1, node2));
        assertThat(summary.getUnavailableReplicas(), containsInAnyOrder(node3));
        verifyReplicaStateObserved(ReplicaState.CONSISTENT, DEFAULT_VERSION.getVersion(), START_TIMESTAMP + 2L);
    }

    @Test
    public void testFailedReplicaIsRemovedFromWriteLockAfterTokenYielded() throws Exception {
        Future<RepositoryConsistencySummary> future = start();

        //Grant tokens for 2 of the 3 nodes
        grantWriteTokens(DEFAULT_VERSION, node1, node3);

        //Mock node3 first asking back its token and then failing (RPC terminates) before the write lock is achieved.
        //If the node is not removed from the write lock, this will cause the write lock to hang waiting for node3's
        //token to be re-granted.
        assertTrue(replicaTransactionCallback.yieldToken(node3));
        replicaTransactionCallback.onReplicaError(node3);

        //Grant the last token and verify that the operation terminates
        grantWriteTokens(DEFAULT_VERSION, node2);

        RepositoryConsistencySummary summary = future.get(TIMEOUT.toMillis(), TimeUnit.MILLISECONDS);
        assertEquals(DEFAULT_VERSION, summary.getConsensusVersion());
        assertThat(summary.getConsistentReplicas(), containsInAnyOrder(node1, node2));
        assertThat(summary.getUnavailableReplicas(), containsInAnyOrder(node3));
        verifyReplicaStateObserved(ReplicaState.CONSISTENT, DEFAULT_VERSION.getVersion(), START_TIMESTAMP + 2L);
    }

    @Test
    public void testLocalReplicaAhead() throws Exception {
        ReplicaVersion localVersion = new ReplicaVersion("other-content", DEFAULT_VERSION.getVersion() + 1L);

        Future<RepositoryConsistencySummary> future = start();

        grantWriteTokens(DEFAULT_VERSION, node1, node3);
        grantWriteTokens(localVersion, node2);

        RepositoryConsistencySummary summary = future.get(TIMEOUT.toMillis(), TimeUnit.MILLISECONDS);
        assertEquals(localVersion, summary.getConsensusVersion());
        assertThat(summary.getInconsistentReplicas(), containsInAnyOrder(node1, node3));
        verifyReplicaStateObserved(ReplicaState.CONSISTENT, localVersion.getVersion(), START_TIMESTAMP + 2L);
    }

    @Test
    public void testLocalReplicaBehind() throws Exception {
        ReplicaVersion localVersion = new ReplicaVersion("other-content", DEFAULT_VERSION.getVersion() - 1L);

        Future<RepositoryConsistencySummary> future = start();

        grantWriteTokens(DEFAULT_VERSION, node1, node3);
        grantWriteTokens(localVersion, node2);

        RepositoryConsistencySummary summary = future.get(TIMEOUT.toMillis(), TimeUnit.MILLISECONDS);
        assertEquals(DEFAULT_VERSION, summary.getConsensusVersion());
        assertThat(summary.getInconsistentReplicas(), Matchers.contains(node2));
        verifyReplicaStateObserved(ReplicaState.INCONSISTENT, DEFAULT_VERSION.getVersion(), START_TIMESTAMP + 2L);
    }

    @Test
    public void testLocalReplicaHasDivergentContentHash() throws Exception {
        ReplicaVersion localVersion = new ReplicaVersion("other-content", DEFAULT_VERSION.getVersion());

        Future<RepositoryConsistencySummary> future = start();

        grantWriteTokens(DEFAULT_VERSION, node1, node3);
        grantWriteTokens(localVersion, node2);

        RepositoryConsistencySummary summary = future.get(TIMEOUT.toMillis(), TimeUnit.MILLISECONDS);
        assertEquals(DEFAULT_VERSION, summary.getConsensusVersion());
        assertThat(summary.getInconsistentReplicas(), Matchers.contains(node2));
        verifyReplicaStateObserved(ReplicaState.INCONSISTENT, DEFAULT_VERSION.getVersion(), START_TIMESTAMP + 2L);
    }

    @Test
    public void testRemoteReplicaAhead() throws Exception {
        ReplicaVersion divergentVersion = new ReplicaVersion("other-content", DEFAULT_VERSION.getVersion() + 1L);

        Future<RepositoryConsistencySummary> future = start();

        grantWriteTokens(DEFAULT_VERSION, node1, node2);
        grantWriteTokens(divergentVersion, node3);

        RepositoryConsistencySummary summary = future.get(TIMEOUT.toMillis(), TimeUnit.MILLISECONDS);
        assertEquals(divergentVersion, summary.getConsensusVersion());
        assertThat(summary.getInconsistentReplicas(), containsInAnyOrder(node1, node2));
        verifyReplicaStateObserved(ReplicaState.INCONSISTENT, divergentVersion.getVersion(), START_TIMESTAMP + 2L);
    }

    @Test
    public void testRemoteReplicaBehind() throws Exception {
        ReplicaVersion divergentVersion = new ReplicaVersion("other-content", DEFAULT_VERSION.getVersion() - 1L);

        Future<RepositoryConsistencySummary> future = start();

        grantWriteTokens(DEFAULT_VERSION, node1, node2);
        grantWriteTokens(divergentVersion, node3);

        RepositoryConsistencySummary summary = future.get(TIMEOUT.toMillis(), TimeUnit.MILLISECONDS);
        assertEquals(DEFAULT_VERSION, summary.getConsensusVersion());
        assertThat(summary.getInconsistentReplicas(), Matchers.contains(node3));
        verifyReplicaStateObserved(ReplicaState.CONSISTENT, DEFAULT_VERSION.getVersion(), START_TIMESTAMP + 2L);
    }

    @Test
    public void testRemoteReplicaHasDivergentContentHash() throws Exception {
        ReplicaVersion divergentVersion = new ReplicaVersion("other-content", DEFAULT_VERSION.getVersion());

        Future<RepositoryConsistencySummary> future = start();

        grantWriteTokens(DEFAULT_VERSION, node1, node2);
        grantWriteTokens(divergentVersion, node3);

        RepositoryConsistencySummary summary = future.get(TIMEOUT.toMillis(), TimeUnit.MILLISECONDS);
        assertEquals(DEFAULT_VERSION, summary.getConsensusVersion());
        assertThat(summary.getInconsistentReplicas(), Matchers.contains(node3));
        verifyReplicaStateObserved(ReplicaState.CONSISTENT, DEFAULT_VERSION.getVersion(), START_TIMESTAMP + 2L);
    }

    @Test
    public void testRemoteReplicaUnavailable() throws Exception {
        Future<RepositoryConsistencySummary> future = start();

        grantWriteTokens(DEFAULT_VERSION, node1, node2);
        mockReplicaUnavailable(node3);

        RepositoryConsistencySummary summary = future.get(TIMEOUT.toMillis(), TimeUnit.MILLISECONDS);
        assertEquals(DEFAULT_VERSION, summary.getConsensusVersion());
        assertThat(summary.getUnavailableReplicas(), Matchers.contains(node3));
        assertThat(summary.getInconsistentReplicas(), Matchers.empty());
        verifyReplicaStateObserved(ReplicaState.CONSISTENT, DEFAULT_VERSION.getVersion(), START_TIMESTAMP + 2L);

        verify(writeLock).remove(node3);
    }

    @Test
    public void testMajorityOfReplicasMissingAtSameVersion() throws Exception {
        Future<RepositoryConsistencySummary> future = start();

        assertFalse(future.isDone());
        //Grant the majority of write tokens with null hash and a valid content hash with the same version
        grantWriteTokens(new ReplicaVersion(RefHasher.NULL_HASH, 1L), node1, node3);
        grantWriteTokens(new ReplicaVersion("content-hash", 1L), node2);

        RepositoryConsistencySummary summary = future.get(TIMEOUT.toMillis(), TimeUnit.MILLISECONDS);

        assertEquals("content-hash", summary.getConsensusVersion().getContentHash());
        assertEquals(1L, summary.getConsensusVersion().getVersion());
        verifyReplicaStateObserved(ReplicaState.CONSISTENT, 1L, START_TIMESTAMP + 2L);
        verify(writeLock).close();
    }

    @Test
    public void testWriteLockError() {
        Future<RepositoryConsistencySummary> future = start();

        waitFor("Never blocked waiting on write lock", () -> writeLockFuture != null, TIMEOUT);
        writeLockFuture.completeExceptionally(new RuntimeException("mocking write lock failure"));

        ExecutionException exception = assertThrows(ExecutionException.class,
                () -> future.get(TIMEOUT.toMillis(), TimeUnit.MILLISECONDS));
        Throwable cause = exception.getCause();
        assertThat(cause, instanceOf(InternalFailureException.class));
        assertEquals("mocking write lock failure", Throwables.getRootCause(cause).getMessage());

        //Verify that the write lock was released, even if we failed to verify consistency
        verify(writeLock).close();

        logListener.assertWarn("Failed to verify replica consistency");
    }

    @Test
    public void testYieldToken() throws Exception {
        Future<RepositoryConsistencySummary> future = start();

        //Simulate node1 and node3 granting their tokens
        grantWriteTokens(DEFAULT_VERSION, node1, node3);
        //Simulate node3 asking for its token back, then re-granting the token at a different version
        assertTrue(requestYieldToken(node3));
        ReplicaVersion node3Version = new ReplicaVersion("other-contenthash", DEFAULT_VERSION.getVersion() + 3L);
        grantWriteTokens(node3Version, node3);
        //And finally node2's token is also granted, and verification can completed
        grantWriteTokens(DEFAULT_VERSION, node2);

        RepositoryConsistencySummary summary = future.get(TIMEOUT.toMillis(), TimeUnit.MILLISECONDS);
        assertEquals(node3Version, summary.getConsensusVersion());
        assertThat(summary.getInconsistentReplicas(), containsInAnyOrder(node1, node2));
        verifyReplicaStateObserved(ReplicaState.INCONSISTENT, node3Version.getVersion(), START_TIMESTAMP + 2L);
    }

    private static void waitFor(String message, BooleanSupplier predicate, Duration timeout) {
        Instant deadline = Instant.now().plus(timeout);
        do {
            if (predicate.getAsBoolean()) {
                return;
            }
            long remaining = Math.max(0L, Duration.between(Instant.now(), deadline).toMillis());
            LockSupport.parkNanos(Math.min(remaining, TimeUnit.MILLISECONDS.toNanos(5L)));
        } while (Instant.now().isBefore(deadline));

        fail(message);
    }

    private void grantWriteTokens(ReplicaVersion replicaVersion, MeshNode... nodes) {
        //Wait for the async thread to request the replica tokens
        waitFor("ReplicaTransactionCallback was never set", () -> replicaTransactionCallback != null, TIMEOUT);

        for (MeshNode node : nodes) {
            writeTokensPending.remove(node);
            if (node.isRemote()) {
                waitFor("Write token has not been requested for " + node,
                        () -> writeTokenFutureByReplica.containsKey(node), TIMEOUT);
                replicaTransactionCallback.onTokenGranted(node, replicaVersion);
                CompletableFuture<Void> tokenFuture = writeTokenFutureByReplica.get(node);
                tokenFuture.complete(null);
            } else {
                when(repositoryHashService.getContentHash(repository)).thenReturn(replicaVersion.getContentHash());
                when(repositoryAttributeManager.getVersion(REPOSITORY_ID)).thenReturn(replicaVersion.getVersion());
            }
        }
        //If all write tokens have been granted, trip the write lock
        maybeCompleteWriteLock();
    }

    private void maybeCompleteWriteLock() {
        if (writeTokensPending.isEmpty()) {
            writeLockFuture.complete(null);
        }
    }

    private void mockReplicaUnavailable(MeshNode node) {
        waitFor("Write token has not been requested for " + node,
                () -> writeTokenFutureByReplica.containsKey(node), TIMEOUT);
        CompletableFuture<Void> tokenFuture = writeTokenFutureByReplica.get(node);
        tokenFuture.completeExceptionally(Status.UNAVAILABLE
                .withDescription("mocking node unavailable")
                .asRuntimeException());
        maybeCompleteWriteLock();
    }

    private boolean requestYieldToken(MeshNode node) {
        //Wait for the async thread to request the replica tokens
        waitFor("ReplicaTransactionCallback was never set", () -> replicaTransactionCallback != null, TIMEOUT);
        waitFor("Write token has not been requested for " + node,
                () -> writeTokenFutureByReplica.containsKey(node), TIMEOUT);

        return replicaTransactionCallback.yieldToken(node);
    }

    private Future<RepositoryConsistencySummary> start() {
        CompletableFuture<RepositoryConsistencySummary> future = new CompletableFuture<>();
        verificationThread = new Thread(() -> {
            try {
                RepositoryConsistencySummary summary = consistencyVerifier.call();
                future.complete(summary);
            } catch (Throwable t) {
                future.completeExceptionally(t);
            }
        });
        verificationThread.start();
        return future;
    }

    private void verifyReplicaStateObserved(ReplicaState state, long version, long timestamp) {
        ArgumentCaptor<ReplicaStateObservedEvent> eventCaptor =
                ArgumentCaptor.forClass(ReplicaStateObservedEvent.class);
        verify(eventPublisher).publish(eventCaptor.capture());

        ReplicaStateObservedEvent event = eventCaptor.getValue();
        assertEquals(REPOSITORY_ID, event.getRepositoryId());
        assertEquals(state, event.getObservation().getState());
        assertEquals(version, event.getObservation().getVersion());
        assertEquals(timestamp, event.getObservation().getLogicalTimestamp());
    }
}