package com.atlassian.bitbucket.mesh.repository;

import com.atlassian.bitbucket.mesh.execution.ExecutionContext;
import com.atlassian.bitbucket.mesh.execution.ExecutionManager;
import com.atlassian.bitbucket.mesh.execution.GrpcExecutionManager;
import com.atlassian.bitbucket.mesh.request.RequestContext;
import com.atlassian.bitbucket.mesh.request.RequestManager;
import com.atlassian.bitbucket.mesh.rpc.v1.RepositoryServiceGrpc;
import com.atlassian.bitbucket.mesh.rpc.v1.TransactionServiceGrpc;
import io.grpc.Context;
import io.grpc.ServerCall;
import org.hamcrest.MatcherAssert;
import org.junit.Before;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.mockito.Mock;
import org.mockito.junit.MockitoJUnitRunner;

import java.io.Closeable;
import java.io.IOException;
import java.io.UncheckedIOException;
import java.util.Optional;
import java.util.concurrent.CompletableFuture;

import static org.hamcrest.Matchers.instanceOf;
import static org.junit.Assert.*;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.Mockito.*;

@RunWith(MockitoJUnitRunner.class)
public class DefaultRpcRepositoryAccessTrackerTest {

    private static final String PARTITION_0 = "0000";
    private static final String PARTITION_1 = "0001";
    private static final String HIERARCHY_1 = "11111111111111111111";
    private static final String HIERARCHY_2 = "22222222222222222222";
    private static final String HIERARCHY_3 = "33333333333333333333";

    private final ExecutionManager executionManager = new GrpcExecutionManager();
    @Mock
    private RequestManager requestManager;
    private boolean rpcMethodIsExcludedFromTracking;

    private DefaultRpcRepositoryAccessTracker tracker;

    @Before
    public void setUp() {
        when(requestManager.getContext()).thenAnswer(invocation -> {
            if (!executionManager.getContext().isPresent()) {
                //No RPC active
                return Optional.empty();
            }

            //Mock an active RPC request context
            RequestContext context = mock(RequestContext.class);
            ServerCall serverCall = mock(ServerCall.class);
            when(serverCall.getMethodDescriptor()).thenAnswer(ignored -> rpcMethodIsExcludedFromTracking
                    ? TransactionServiceGrpc.getTransactMethod() // excluded from tracking
                    : RepositoryServiceGrpc.getIsEmptyMethod());
            when(context.getRawRequest()).thenReturn(serverCall);
            return Optional.of(context);
        });
        tracker = new DefaultRpcRepositoryAccessTracker(executionManager, requestManager);
    }

    @Test
    public void testDoesNotTrackAccessesFromExcludedRpcMethods() {
        ManagedRepository repository = repository(PARTITION_0, HIERARCHY_1);
        executionManager.withContext(ctx -> {
            //Triggers the mock request context to return an excluded RPC method
            rpcMethodIsExcludedFromTracking = true;

            assertEquals(0, tracker.getActiveRequestCount());
            assertEquals(0, tracker.getActiveRequestCountByPartition(PARTITION_0));
            assertEquals(0, tracker.getActiveRequestCountByHierarchy(PARTITION_0, HIERARCHY_1));

            tracker.onRepositoryAccessed(repository);

            //Verify that the repository access is ignored for all counts counts
            assertEquals(0, tracker.getActiveRequestCount());
            assertEquals(0, tracker.getActiveRequestCountByPartition(PARTITION_0));
            assertEquals(0, tracker.getActiveRequestCountByHierarchy(PARTITION_0, HIERARCHY_1));

            //Verify that draining returns immediately
            assertTrue(tracker.drainAll().toCompletableFuture().isDone());
            assertTrue(tracker.drainHierarchy(PARTITION_0, HIERARCHY_1).toCompletableFuture().isDone());
            assertTrue(tracker.drainPartition(PARTITION_0).toCompletableFuture().isDone());
            return null;
        });
    }

    @Test
    public void testDoesNotTrackAccessesFromNonRpcContexts() {
        assertEquals(0, tracker.getActiveRequestCount());
        assertEquals(0, tracker.getActiveRequestCountByPartition(PARTITION_0));

        Repository repository = repository(PARTITION_0, HIERARCHY_1);

        //No ExecutionContext is active, so the access should not be registered or tracked at all
        tracker.onRepositoryAccessed(repository);

        assertEquals(0, tracker.getActiveRequestCount());
        assertEquals(0, tracker.getActiveRequestCountByPartition(PARTITION_0));

        //Draining all should immediately succeed
        assertTrue(tracker.drainAll().toCompletableFuture().isDone());
    }

    @Test
    public void testDrainingMultipleRpcs() {
        //To mock multiple concurrent RPC requests, we create nested ExecutionContexts in this test. Each context
        //corresponds to a separate RPC request. In order to control when these RPC contexts close, we create them
        //from an open gRPC context.
        Context.ROOT.withCancellation().run(() -> {
            //We've opened a new gRPC context. Any ExecutionContexts we create from here won't be (automatically) closed
            //until the gRPC context is closed.
            ExecutionContext rpc1 = executionManager.withContext(ctx -> ctx);
            ExecutionContext rpc2 = executionManager.withContext(ctx -> ctx);
            ExecutionContext rpc3 = executionManager.withContext(ctx -> ctx);

            Repository repository0_1 = repository(PARTITION_0, HIERARCHY_1);
            Repository repository0_3 = repository(PARTITION_0, HIERARCHY_3);
            Repository repository1_2 = repository(PARTITION_1, HIERARCHY_2);

            //Mock a repository access from rpc1
            rpc1.run(() -> tracker.onRepositoryAccessed(repository0_1));
            assertEquals(1, tracker.getActiveRequestCount());
            assertEquals(1, tracker.getActiveRequestCountByHierarchy(PARTITION_0, HIERARCHY_1));
            assertEquals(0, tracker.getActiveRequestCountByHierarchy(PARTITION_1, HIERARCHY_2));
            assertEquals(0, tracker.getActiveRequestCountByHierarchy(PARTITION_0, HIERARCHY_3));
            assertEquals(1, tracker.getActiveRequestCountByPartition(PARTITION_0));
            assertEquals(0, tracker.getActiveRequestCountByPartition(PARTITION_1));

            //Start draining
            CompletableFuture<Void> drainAllFuture = new CompletableFuture<>();
            CompletableFuture<Void> drainPartition0Future = new CompletableFuture<>();
            CompletableFuture<Void> drainHierarchy1Future = new CompletableFuture<>();
            tracker.drainAll().thenRun(() -> drainAllFuture.complete(null));
            tracker.drainHierarchy(PARTITION_0, HIERARCHY_1).thenRun(() -> drainHierarchy1Future.complete(null));
            tracker.drainPartition(PARTITION_0).thenRun(() -> drainPartition0Future.complete(null));

            assertFalse(drainAllFuture.isDone());
            assertFalse(drainHierarchy1Future.isDone());
            assertFalse(drainPartition0Future.isDone());

            //Mock a rpc2 accessing both repositories
            rpc2.run(() -> tracker.onRepositoryAccessed(repository1_2));
            rpc2.run(() -> tracker.onRepositoryAccessed(repository0_1));

            assertEquals(2, tracker.getActiveRequestCount());
            assertEquals(2, tracker.getActiveRequestCountByHierarchy(PARTITION_0, HIERARCHY_1));
            assertEquals(1, tracker.getActiveRequestCountByHierarchy(PARTITION_1, HIERARCHY_2));
            assertEquals(0, tracker.getActiveRequestCountByHierarchy(PARTITION_0, HIERARCHY_3));
            assertEquals(2, tracker.getActiveRequestCountByPartition(PARTITION_0));
            assertEquals(1, tracker.getActiveRequestCountByPartition(PARTITION_1));
            assertFalse(drainAllFuture.isDone());
            assertFalse(drainHierarchy1Future.isDone());
            assertFalse(drainPartition0Future.isDone());

            //Mock rpc1 finishing
            closeRpc(rpc1);

            assertEquals(1, tracker.getActiveRequestCount());
            assertEquals(1, tracker.getActiveRequestCountByHierarchy(PARTITION_0, HIERARCHY_1));
            assertEquals(1, tracker.getActiveRequestCountByHierarchy(PARTITION_1, HIERARCHY_2));
            assertEquals(0, tracker.getActiveRequestCountByHierarchy(PARTITION_0, HIERARCHY_3));
            assertEquals(1, tracker.getActiveRequestCountByPartition(PARTITION_0));
            assertEquals(1, tracker.getActiveRequestCountByPartition(PARTITION_1));
            assertFalse(drainAllFuture.isDone());
            assertFalse(drainHierarchy1Future.isDone());
            assertFalse(drainPartition0Future.isDone());

            //Mock rpc3 accessing another repository in on partition 0
            rpc3.run(() -> tracker.onRepositoryAccessed(repository0_3));

            assertEquals(2, tracker.getActiveRequestCount());
            assertEquals(1, tracker.getActiveRequestCountByHierarchy(PARTITION_0, HIERARCHY_1));
            assertEquals(1, tracker.getActiveRequestCountByHierarchy(PARTITION_1, HIERARCHY_2));
            assertEquals(1, tracker.getActiveRequestCountByHierarchy(PARTITION_0, HIERARCHY_3));
            assertEquals(2, tracker.getActiveRequestCountByPartition(PARTITION_0));
            assertEquals(1, tracker.getActiveRequestCountByPartition(PARTITION_1));
            assertFalse(drainAllFuture.isDone());
            assertFalse(drainHierarchy1Future.isDone());
            assertFalse(drainPartition0Future.isDone());

            //Mock rpc2 finishing too
            closeRpc(rpc2);

            assertEquals(1, tracker.getActiveRequestCount());
            assertEquals(0, tracker.getActiveRequestCountByHierarchy(PARTITION_0, HIERARCHY_1));
            assertEquals(0, tracker.getActiveRequestCountByHierarchy(PARTITION_1, HIERARCHY_2));
            assertEquals(1, tracker.getActiveRequestCountByHierarchy(PARTITION_0, HIERARCHY_3));
            assertEquals(1, tracker.getActiveRequestCountByPartition(PARTITION_0));
            assertEquals(0, tracker.getActiveRequestCountByPartition(PARTITION_1));
            assertFalse(drainAllFuture.isDone());
            assertTrue(drainHierarchy1Future.isDone());
            assertFalse(drainPartition0Future.isDone());

            //And finally finish rpc3 too
            closeRpc(rpc3);

            assertEquals(0, tracker.getActiveRequestCount());
            assertEquals(0, tracker.getActiveRequestCountByHierarchy(PARTITION_0, HIERARCHY_1));
            assertEquals(0, tracker.getActiveRequestCountByHierarchy(PARTITION_1, HIERARCHY_2));
            assertEquals(0, tracker.getActiveRequestCountByHierarchy(PARTITION_0, HIERARCHY_3));
            assertEquals(0, tracker.getActiveRequestCountByPartition(PARTITION_0));
            assertEquals(0, tracker.getActiveRequestCountByPartition(PARTITION_1));
            assertTrue(drainAllFuture.isDone());
            assertTrue(drainPartition0Future.isDone());
        });
    }

    @Test
    public void testDrainMultipleTimes() {
        //To mock multiple RPC requests, we create nested ExecutionContexts in this test. Each context corresponds
        //to a separate RPC request. In order to control when these RPC contexts close, we create them from an open
        //gRPC context.
        Context.ROOT.withCancellation().run(() -> {
            //Loop 3 times to validate that draining works as expected after a previous drain command completed
            for (int i = 0; i < 3; ++i) {
                //We've opened a new gRPC context. Any ExecutionContexts we create from here won't be (automatically)
                //closed until the gRPC context is closed.
                ExecutionContext rpc = executionManager.withContext(ctx -> ctx);

                CompletableFuture<Void> drainFuture = new CompletableFuture<>();
                rpc.run(() -> tracker.onRepositoryAccessed(repository(PARTITION_0, HIERARCHY_1)));

                tracker.drainAll().thenRun(() -> drainFuture.complete(null));
                assertFalse(drainFuture.isDone());

                closeRpc(rpc);
                assertTrue(drainFuture.isDone());
            }
        });
    }

    @Test
    public void testNonManagedRepositoryAccess() {
        Repository bitbucketRepository = bitbucketRepository(HIERARCHY_2);

        executionManager.withContext(ctx -> {
            tracker.onRepositoryAccessed(bitbucketRepository);

            //Only the 'all' tracker should be updated because 'classic' repositories don't have a partition
            assertEquals(1, tracker.getActiveRequestCount());

            return null;
        });
        assertEquals(0, tracker.getActiveRequestCount());
    }

    @Test
    public void testTracksActiveRpcRequests() {
        Repository repository0_1 = repository(PARTITION_0, HIERARCHY_1);
        Repository repository0_3 = repository(PARTITION_0, HIERARCHY_3);
        Repository repository1_2 = repository(PARTITION_1, HIERARCHY_2);

        CompletableFuture<Void> drainFuture = new CompletableFuture<>();
        executionManager.withContext(ctx -> {
            tracker.onRepositoryAccessed(repository0_1);

            //Start a drain operation. We don't expect this one to complete until the RPC completes
            tracker.drainAll().thenRun(() -> drainFuture.complete(null));
            assertFalse(drainFuture.isDone());

            //Verify that the repository access increases the counts for all, the partition and the hierarchy
            assertEquals(1, tracker.getActiveRequestCount());
            assertEquals(1, tracker.getActiveRequestCountByPartition(PARTITION_0));
            assertEquals(1, tracker.getActiveRequestCountByHierarchy(PARTITION_0, HIERARCHY_1));
            //Other partition counts should be unaffected
            assertEquals(0, tracker.getActiveRequestCountByPartition(PARTITION_1));
            assertEquals(0, tracker.getActiveRequestCountByHierarchy(PARTITION_1, HIERARCHY_2));

            //Access another repository on another partition from the same RPC
            tracker.onRepositoryAccessed(repository1_2);

            //Total active request count should not have gone up
            assertEquals(1, tracker.getActiveRequestCount());
            assertEquals(1, tracker.getActiveRequestCountByPartition(PARTITION_0));
            assertEquals(1, tracker.getActiveRequestCountByHierarchy(PARTITION_0, HIERARCHY_1));
            assertEquals(1, tracker.getActiveRequestCountByPartition(PARTITION_1));
            assertEquals(1, tracker.getActiveRequestCountByHierarchy(PARTITION_1, HIERARCHY_2));

            //Access another repository on partition 0
            tracker.onRepositoryAccessed(repository0_3);

            assertEquals(1, tracker.getActiveRequestCount());
            assertEquals(1, tracker.getActiveRequestCountByPartition(PARTITION_0));
            assertEquals(1, tracker.getActiveRequestCountByHierarchy(PARTITION_0, HIERARCHY_1));
            assertEquals(1, tracker.getActiveRequestCountByHierarchy(PARTITION_0, HIERARCHY_3));
            assertEquals(1, tracker.getActiveRequestCountByPartition(PARTITION_1));
            assertEquals(1, tracker.getActiveRequestCountByHierarchy(PARTITION_1, HIERARCHY_2));

            assertFalse(drainFuture.isDone());
            return null;
        });

        assertEquals(0, tracker.getActiveRequestCount());
        assertEquals(0, tracker.getActiveRequestCountByPartition(PARTITION_0));
        assertEquals(0, tracker.getActiveRequestCountByHierarchy(PARTITION_0, HIERARCHY_1));
        assertEquals(0, tracker.getActiveRequestCountByHierarchy(PARTITION_0, HIERARCHY_3));
        assertEquals(0, tracker.getActiveRequestCountByPartition(PARTITION_1));
        assertEquals(0, tracker.getActiveRequestCountByHierarchy(PARTITION_1, HIERARCHY_2));

        //Draining should have completed now
        assertTrue(drainFuture.isDone());
    }

    private static BitbucketServerRepository bitbucketRepository(String hierarchy) {
        BitbucketServerRepository repository = mock(BitbucketServerRepository.class);
        lenient().when(repository.getHierarchyId()).thenReturn(hierarchy);
        when(repository.accept(any())).thenAnswer(invocation -> {
            RepositoryVisitor visitor = invocation.getArgument(0);
            return visitor.visit(repository);
        });
        return repository;
    }

    private static void closeRpc(ExecutionContext executionContext) {
        MatcherAssert.assertThat(executionContext, instanceOf(Closeable.class));
        try {
            ((Closeable) executionContext).close();
        } catch (IOException e) {
            throw new UncheckedIOException(e);
        }
    }

    private static ManagedRepository repository(String partition, String hierarchy) {
        ManagedRepository repository = mock(ManagedRepository.class);
        when(repository.getPartition()).thenReturn(partition);
        when(repository.getHierarchyId()).thenReturn(hierarchy);
        when(repository.accept(any())).thenAnswer(invocation -> {
            RepositoryVisitor visitor = invocation.getArgument(0);
            return visitor.visit(repository);
        });
        return repository;
    }
}