package com.atlassian.bitbucket.mesh.client;

import com.atlassian.bitbucket.mesh.coordination.LogicalClock;
import com.atlassian.bitbucket.mesh.network.MeshNode;
import com.atlassian.bitbucket.mesh.rpc.v1.*;
import com.atlassian.bitbucket.mesh.transaction.ReplicaTransactionCallback;
import com.atlassian.bitbucket.mesh.transaction.ReplicaVersion;
import io.grpc.Status;
import io.grpc.stub.ClientCallStreamObserver;
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.CompletionStage;
import java.util.function.BiFunction;

import static org.junit.Assert.*;
import static org.mockito.Mockito.*;

@RunWith(MockitoJUnitRunner.class)
public class ReplicaConsistencyResponseObserverTest {

    private static final String REPOSITORY_ID = "p/0000/h/123456789012345678980/r/82374";
    private static final String TRANSACTION_ID = "tx-id";
    @Mock
    private ReplicaTransactionCallback callback;
    @Mock
    private ClientCallStreamObserver<RpcVerifyConsistencyLeaderFragment> clientStreamObserver;
    @Captor
    private ArgumentCaptor<RpcVerifyConsistencyLeaderFragment> fragmentCaptor;
    @Mock
    private MeshNode localNode;
    @Mock
    private LogicalClock logicalClock;
    @Mock
    private MeshNode replicaNode;
    @Captor
    private ArgumentCaptor<ReplicaVersion> replicaVersionCaptor;
    private ReplicaConsistencyResponseObserver responseObserver;
    @Mock
    private BiFunction<Void, Throwable, Void> tokenStageHandler;

    @Before
    public void setUp() {
        responseObserver = new ReplicaConsistencyResponseObserver(logicalClock, replicaNode,
                new ReplicateParameters.Builder(localNode, TRANSACTION_ID).callback(callback).build());
        responseObserver.beforeStart(clientStreamObserver);
    }

    @Test
    public void testCancel() {
        //Request the token, and then cancel the RPC
        CompletionStage<Void> tokenStage = responseObserver.requestTokenAndReplicaVersion(REPOSITORY_ID, 17L);
        tokenStage.handle(tokenStageHandler);
        responseObserver.cancel("testing cancellation");

        //Verify that the RPC is cancelled
        verify(clientStreamObserver).onNext(any());
        verify(clientStreamObserver).cancel("testing cancellation", null);
        verifyNoInteractions(tokenStageHandler);

        //Simulate onError being called with the CANCELLED status
        responseObserver.onError(Status.CANCELLED.asRuntimeException());
        verify(tokenStageHandler).apply(eq(null), any());
    }

    @Test
    public void testHappyFlow() {
        verifyNoInteractions(clientStreamObserver);
        assertEquals(replicaNode, responseObserver.getReplica());

        //Request the token
        CompletionStage<Void> tokenStage = responseObserver.requestTokenAndReplicaVersion(REPOSITORY_ID, 17L);

        //Verify that the token request is sent to the replica
        assertNotNull(tokenStage);
        verify(clientStreamObserver).onNext(fragmentCaptor.capture());
        verifyNoMoreInteractions(clientStreamObserver);

        RpcVerifyConsistencyLeaderFragment fragment = fragmentCaptor.getValue();
        assertEquals(RpcVerifyConsistencyLeaderFragment.newBuilder()
                        .setLogicalTimestamp(17L)
                        .setTokenRequest(RpcWriteTokenRequest.newBuilder()
                                .setTransactionId(TRANSACTION_ID)
                                .setRepository(REPOSITORY_ID))
                        .build(),
                fragment);

        //Verify that the tokenStage hasn't been completed yet
        tokenStage.handle(tokenStageHandler);
        verifyNoInteractions(tokenStageHandler);

        //Simulate the replica granting its token
        responseObserver.onNext(RpcVerifyConsistencyFollowerFragment.newBuilder()
                .setLogicalTimestamp(21L)
                .setTokenGranted(RpcWriteTokenGranted.newBuilder()
                        .setReplicaVersion(RpcReplicaVersion.newBuilder()
                                .setContentHash("content")
                                .setVersion(99L))
                        .build())
                .build());
        verify(logicalClock).onReceived(21L);

        //Verify that the callback is invoked and that the token stage handler is completed
        verify(callback).onTokenGranted(eq(replicaNode), replicaVersionCaptor.capture());
        verify(tokenStageHandler).apply(any(), eq(null));
        ReplicaVersion replicaVersion = replicaVersionCaptor.getValue();
        assertEquals("content", replicaVersion.getContentHash());
        assertEquals(99L, replicaVersion.getVersion());

        //Simulate the leader sending a summary
        RpcRepositoryConsistencySummary summary = RpcRepositoryConsistencySummary.newBuilder()
                .addInconsistentReplicas("node-1")
                .addUnavailableReplicas("node-2")
                .setContentHash("content-2")
                .setVersion(101L)
                .build();
        responseObserver.summarize(summary, 23L);

        //Verify that the summary is forwarded and that the RPC is completed
        verify(clientStreamObserver, times(2)).onNext(fragmentCaptor.capture());
        verify(clientStreamObserver).onCompleted();

        fragment = fragmentCaptor.getValue();
        assertEquals(23L, fragment.getLogicalTimestamp());
        assertEquals(summary, fragment.getConsistencySummary());

        verifyNoMoreInteractions(clientStreamObserver);

        //Simulate responseObserver completing too
        responseObserver.onCompleted();
    }

    @Test
    public void testTokenGrantedThenYieldAndGrantedAgain() {
        //Request the token
        CompletionStage<Void> tokenStage = responseObserver.requestTokenAndReplicaVersion(REPOSITORY_ID, 17L);

        //Verify that the token has not been granted yet
        tokenStage.handle(tokenStageHandler);
        verifyNoInteractions(tokenStageHandler, callback);

        //Simulate the replica granting its token
        responseObserver.onNext(RpcVerifyConsistencyFollowerFragment.newBuilder()
                .setLogicalTimestamp(21L)
                .setTokenGranted(RpcWriteTokenGranted.newBuilder()
                        .setReplicaVersion(RpcReplicaVersion.newBuilder()
                                .setContentHash("content-1")
                                .setVersion(99L))
                        .build())
                .build());
        verify(clientStreamObserver).onNext(any());
        verify(tokenStageHandler).apply(null, null);
        verify(callback).onTokenGranted(eq(replicaNode), replicaVersionCaptor.capture());

        //Simulate the replica requesting its token back, and simulate that the token can be yielded
        when(callback.yieldToken(eq(replicaNode))).thenReturn(true);
        when(logicalClock.tick()).thenReturn(28L);
        responseObserver.onNext(RpcVerifyConsistencyFollowerFragment.newBuilder()
                .setLogicalTimestamp(23L)
                .setYieldTokenRequest(RpcYieldWriteTokenRequest.getDefaultInstance())
                .build());
        //Verify that the callback is forwarded and that the token was yielded
        verify(callback).yieldToken(replicaNode);
        verify(clientStreamObserver, times(2)).onNext(fragmentCaptor.capture());
        verify(logicalClock).tick();

        RpcVerifyConsistencyLeaderFragment fragment = fragmentCaptor.getValue();
        assertTrue(fragment.hasTokenYielded());
        assertEquals(28L, fragment.getLogicalTimestamp());

        //Simulate the token being granted again, but with different hashes and version
        responseObserver.onNext(RpcVerifyConsistencyFollowerFragment.newBuilder()
                .setLogicalTimestamp(29L)
                .setTokenGranted(RpcWriteTokenGranted.newBuilder()
                        .setReplicaVersion(RpcReplicaVersion.newBuilder()
                                .setContentHash("content-2")
                                .setVersion(101L))
                        .build())
                .build());

        //Verify that the newest replica version is passed on to the callback
        verify(callback, times(2)).onTokenGranted(eq(replicaNode), replicaVersionCaptor.capture());
        ReplicaVersion replicaVersion = replicaVersionCaptor.getValue();
        assertEquals("content-2", replicaVersion.getContentHash());
        assertEquals(101L, replicaVersion.getVersion());
    }

    @Test
    public void testYieldRequestRefused() {
        //Request the token
        CompletionStage<Void> tokenStage = responseObserver.requestTokenAndReplicaVersion(REPOSITORY_ID, 17L);

        //Verify that the token has not been granted yet
        tokenStage.handle(tokenStageHandler);
        verifyNoInteractions(tokenStageHandler, callback);

        //Simulate the replica granting its token
        responseObserver.onNext(RpcVerifyConsistencyFollowerFragment.newBuilder()
                .setLogicalTimestamp(21L)
                .setTokenGranted(RpcWriteTokenGranted.newBuilder()
                        .setReplicaVersion(RpcReplicaVersion.newBuilder()
                                .setContentHash("content-1")
                                .setVersion(99L))
                        .build())
                .build());
        verify(clientStreamObserver).onNext(any());
        verify(tokenStageHandler).apply(null, null);
        verify(callback).onTokenGranted(eq(replicaNode), replicaVersionCaptor.capture());

        //Simulate the replica requesting its token back, and simulate that the token cannot be yielded
        responseObserver.onNext(RpcVerifyConsistencyFollowerFragment.newBuilder()
                .setLogicalTimestamp(23L)
                .setYieldTokenRequest(RpcYieldWriteTokenRequest.getDefaultInstance())
                .build());
        //Verify that the callback is forwarded and that the token was yielded
        verify(callback).yieldToken(replicaNode);
        verifyNoMoreInteractions(clientStreamObserver);
        verify(logicalClock, never()).tick();
    }
}