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.ExecutionContext;
import com.atlassian.bitbucket.mesh.execution.ExecutionManager;
import com.atlassian.bitbucket.mesh.network.MeshManager;
import com.atlassian.bitbucket.mesh.network.MeshNode;
import com.atlassian.bitbucket.mesh.repository.*;
import com.atlassian.bitbucket.mesh.rpc.v1.*;
import io.grpc.Status;
import io.grpc.stub.StreamObserver;
import org.junit.Before;
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.util.concurrent.atomic.AtomicLong;

import static com.atlassian.bitbucket.mesh.repository.ReplicaState.CONSISTENT;
import static com.atlassian.bitbucket.mesh.repository.ReplicaState.INCONSISTENT;
import static java.util.Optional.of;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertTrue;
import static org.mockito.ArgumentMatchers.*;
import static org.mockito.Mockito.*;

@RunWith(MockitoJUnitRunner.class)
public class VerifyConsistencyFollowerTest {

    private static final String HASH_CONTENT = "content";
    private static final String NODE_ID = "node-id";
    private static final String REPOSITORY_ID = "repository-id";
    private static final String TX_ID = "tx-id";
    private static final long VERSION = 128L;

    @Captor
    private ArgumentCaptor<ReplicaStateObservedEvent> eventCaptor;
    @Mock
    private EventPublisher eventPublisher;
    @Mock
    private ExecutionContext executionContext;
    @Mock
    private ExecutionManager executionManager;
    @Mock
    private StreamObserver<RpcVerifyConsistencyFollowerFragment> followerObserver;
    @Captor
    private ArgumentCaptor<RpcVerifyConsistencyFollowerFragment> fragmentCaptor;
    @Mock
    private MeshNode localNode;
    @Mock
    private LogicalClock logicalClock;
    @Mock
    private MeshManager meshManager;
    @Mock
    private Repository repository;
    @Mock
    private RepositoryAttributeManager repositoryAttributeManager;
    @Mock
    private RepositoryHashService repositoryHashService;
    @Mock
    private RepositoryManager repositoryManager;
    @Captor
    private ArgumentCaptor<WriteTokenHandler> tokenHandlerCaptor;
    private VerifyConsistencyFollower verifyConsistencyFollower;
    @Mock
    private WriteToken writeToken;
    @Mock
    private WriteToken.Promise writeTokenPromise;
    @Mock
    private WriteTokenRegistry writeTokenRegistry;

    @Before
    public void setUp() {
        when(executionManager.getContext()).thenReturn(of(executionContext));
        when(repositoryManager.findById(REPOSITORY_ID, true)).thenReturn(repository);
        when(repository.getId()).thenReturn(REPOSITORY_ID);
        when(repository.exists()).thenReturn(true);
        when(writeTokenRegistry.acquireToken(eq(repository), eq(TX_ID), anyLong(), any()))
                .thenReturn(writeTokenPromise);
        AtomicLong timestamp = new AtomicLong(12L);
        when(logicalClock.tick()).thenAnswer(ignored -> timestamp.incrementAndGet());
        when(logicalClock.onReceived(anyLong())).thenAnswer(invocation -> {
            long value = invocation.getArgument(0);
            return timestamp.updateAndGet(curValue -> Math.max(curValue, value));
        });
        when(meshManager.getLocalNode()).thenReturn(localNode);
        when(localNode.getId()).thenReturn(NODE_ID);
        when(repositoryHashService.getContentHash(repository)).thenReturn(HASH_CONTENT);
        when(repositoryAttributeManager.getVersion(REPOSITORY_ID)).thenReturn(VERSION);

        verifyConsistencyFollower = new VerifyConsistencyFollower(followerObserver, eventPublisher, executionManager,
                logicalClock, meshManager, repositoryAttributeManager, repositoryHashService, repositoryManager,
                writeTokenRegistry);
    }

    @Test
    public void testHappyFlowConsistent() {
        //Simulate receiving a token request from the leader
        verifyConsistencyFollower.onNext(RpcVerifyConsistencyLeaderFragment.newBuilder()
                .setLogicalTimestamp(9L)
                .setTokenRequest(RpcWriteTokenRequest.newBuilder()
                        .setRepository(REPOSITORY_ID)
                        .setTransactionId(TX_ID))
                .build());

        //Verify that the write token has been requested for the local replica, but not sent to the leader yet
        verify(logicalClock).onReceived(9L);
        verify(repositoryHashService).resetContentHash(repository);
        verify(writeTokenRegistry).acquireToken(eq(repository), eq(TX_ID), eq(9L), tokenHandlerCaptor.capture());
        verifyNoInteractions(followerObserver);

        //Simulate the token being granted
        WriteTokenHandler handler = tokenHandlerCaptor.getValue();
        handler.onGranted(writeToken);

        //Verify that an RpcWriteTokenGranted message is sent to the leader
        verify(followerObserver).onNext(fragmentCaptor.capture());
        RpcVerifyConsistencyFollowerFragment fragment = fragmentCaptor.getValue();
        //Clock should have ticked once when sending the fragment
        verify(logicalClock).tick();
        assertEquals(13L, fragment.getLogicalTimestamp());
        assertTrue(fragment.hasTokenGranted());
        //The local hashes and version should be sent to the leader
        RpcReplicaVersion replicaVersion = fragment.getTokenGranted().getReplicaVersion();
        assertEquals(HASH_CONTENT, replicaVersion.getContentHash());
        assertEquals(VERSION, replicaVersion.getVersion());

        //Simulate the summary being received and this node being consistent
        verifyConsistencyFollower.onNext(RpcVerifyConsistencyLeaderFragment.newBuilder()
                .setLogicalTimestamp(14L)
                .setConsistencySummary(RpcRepositoryConsistencySummary.newBuilder()
                        .addConsistentReplicas(NODE_ID)
                        .setContentHash(HASH_CONTENT)
                        .setVersion(VERSION))
                .build());

        //Verify that a CONSISTENT ReplicaStateEvent is raised
        verify(eventPublisher).publish(eventCaptor.capture());
        ReplicaStateObservedEvent event = eventCaptor.getValue();
        assertEquals(REPOSITORY_ID, event.getRepositoryId());
        assertEquals(CONSISTENT, event.getObservation().getState());
        assertEquals(14L, event.getObservation().getLogicalTimestamp());
        assertEquals(VERSION, event.getObservation().getVersion());
        assertTrue(event.getObservation().isForce());

        //Verify that the token is released
        verify(writeTokenPromise).close();

        //Simulate the RPC being completed
        verifyConsistencyFollower.onCompleted();
        verify(followerObserver).onCompleted();
        verifyNoMoreInteractions(writeTokenPromise, eventPublisher, followerObserver);
    }

    @Test
    public void testHappyFlowInconsistent() {
        //Simulate receiving a token request from the leader
        verifyConsistencyFollower.onNext(RpcVerifyConsistencyLeaderFragment.newBuilder()
                .setLogicalTimestamp(9L)
                .setTokenRequest(RpcWriteTokenRequest.newBuilder()
                        .setRepository(REPOSITORY_ID)
                        .setTransactionId(TX_ID))
                .build());

        //Verify that the write token has been requested for the local replica, but not sent to the leader yet
        verify(logicalClock).onReceived(9L);
        verify(writeTokenRegistry).acquireToken(eq(repository), eq(TX_ID), eq(9L), tokenHandlerCaptor.capture());
        verifyNoInteractions(followerObserver);

        //Simulate the token being granted
        WriteTokenHandler handler = tokenHandlerCaptor.getValue();
        handler.onGranted(writeToken);

        //Verify that an RpcWriteTokenGranted message is sent to the leader
        verify(followerObserver).onNext(any());

        //Simulate the summary being received and this node being inconsistent
        verifyConsistencyFollower.onNext(RpcVerifyConsistencyLeaderFragment.newBuilder()
                .setLogicalTimestamp(14L)
                .setConsistencySummary(RpcRepositoryConsistencySummary.newBuilder()
                        .addInconsistentReplicas(NODE_ID)
                        .setContentHash("other")
                        .setVersion(VERSION + 1L))
                .build());

        //Verify that a INCONSISTENT ReplicaStateEvent is raised
        verify(eventPublisher).publish(eventCaptor.capture());
        ReplicaStateObservedEvent event = eventCaptor.getValue();
        assertEquals(REPOSITORY_ID, event.getRepositoryId());
        assertEquals(INCONSISTENT, event.getObservation().getState());
        assertEquals(14L, event.getObservation().getLogicalTimestamp());
        assertEquals(VERSION + 1L, event.getObservation().getVersion());
        assertTrue(event.getObservation().isForce());

        //Verify that the token is released
        verify(writeTokenPromise).close();

        //Simulate the RPC being completed
        verifyConsistencyFollower.onCompleted();
        verify(followerObserver).onCompleted();
        verifyNoMoreInteractions(writeTokenPromise, eventPublisher, followerObserver);
    }

    @Test
    public void testTokenIsReleasedOnLeaderError() {
        //Simulate receiving a token request from the leader
        verifyConsistencyFollower.onNext(RpcVerifyConsistencyLeaderFragment.newBuilder()
                .setLogicalTimestamp(9L)
                .setTokenRequest(RpcWriteTokenRequest.newBuilder()
                        .setRepository(REPOSITORY_ID)
                        .setTransactionId(TX_ID))
                .build());

        //Simulate granting the token
        verify(writeTokenRegistry).acquireToken(eq(repository), eq(TX_ID), eq(9L), tokenHandlerCaptor.capture());
        WriteTokenHandler handler = tokenHandlerCaptor.getValue();
        handler.onGranted(writeToken);

        //Verify that the granted message is sent to the leader
        verify(followerObserver).onNext(any());

        //Simulate the leader running into an error
        verifyConsistencyFollower.onError(Status.INTERNAL.asRuntimeException());

        verify(writeTokenPromise).close();
    }

    @Test
    public void testYieldToken() {
        //Simulate receiving a token request from the leader
        verifyConsistencyFollower.onNext(RpcVerifyConsistencyLeaderFragment.newBuilder()
                .setLogicalTimestamp(9L)
                .setTokenRequest(RpcWriteTokenRequest.newBuilder()
                        .setRepository(REPOSITORY_ID)
                        .setTransactionId(TX_ID))
                .build());

        //Simulate granting the token
        verify(writeTokenRegistry).acquireToken(eq(repository), eq(TX_ID), eq(9L), tokenHandlerCaptor.capture());
        verifyNoInteractions(followerObserver);
        WriteTokenHandler handler = tokenHandlerCaptor.getValue();
        handler.onGranted(writeToken);

        //Verify that the granted message is sent to the leader
        verify(followerObserver).onNext(any());

        //Now simulate the local write lock registry requesting us to yield the token
        handler.yieldToken();

        //Verify that the request to yield has been forwarded to the leader
        verify(followerObserver, times(2)).onNext(fragmentCaptor.capture());
        RpcVerifyConsistencyFollowerFragment fragment = fragmentCaptor.getValue();
        verify(logicalClock, times(2)).tick(); // clock should be ticked on each outgoing message
        assertTrue(fragment.hasYieldTokenRequest());
        //Verify that the token has not been yielded yet
        verify(writeToken, never()).yieldToken();

        //Simulate the leader yielding the token
        verifyConsistencyFollower.onNext(RpcVerifyConsistencyLeaderFragment.newBuilder()
                .setLogicalTimestamp(18L)
                .setTokenYielded(RpcWriteTokenYielded.getDefaultInstance())
                .build());

        //Verify that the token is yielded back to the registry
        verify(logicalClock).onReceived(18L); // clock should be updated
        verify(writeToken).yieldToken();

        //Simulate a concurrent (but earlier) write completing, and the token then being granted again
        when(repositoryHashService.getContentHash(repository)).thenReturn("content-2");
        when(repositoryAttributeManager.getVersion(REPOSITORY_ID)).thenReturn(VERSION + 1L);
        handler.onGranted(writeToken);

        //Verify that the token granted notification has been forwarded to the leader
        verify(followerObserver, times(3)).onNext(fragmentCaptor.capture());
        fragment = fragmentCaptor.getValue();
        verify(logicalClock, times(3)).tick(); // clock should be ticked on each outgoing message
        assertTrue(fragment.hasTokenGranted());
        RpcReplicaVersion replicaVersion = fragment.getTokenGranted().getReplicaVersion();
        assertEquals("content-2", replicaVersion.getContentHash());
        assertEquals(VERSION + 1L, replicaVersion.getVersion());
    }
}