package com.atlassian.bitbucket.mesh.management;

import com.atlassian.bitbucket.mesh.git.exception.FailedPreconditionException;
import com.atlassian.bitbucket.mesh.grpc.GrpcServerShutdownEvent;
import com.atlassian.bitbucket.mesh.junit.LogListener;
import com.atlassian.bitbucket.mesh.request.RequestContext;
import com.atlassian.bitbucket.mesh.request.RequestManager;
import com.atlassian.bitbucket.mesh.rpc.v1.RpcHeartbeat;
import com.atlassian.bitbucket.mesh.rpc.v1.RpcRepositoryReplicaObservation;
import com.atlassian.bitbucket.mesh.rpc.v1.RpcSidebandFollowerFragment;
import com.atlassian.bitbucket.mesh.rpc.v1.RpcSidebandLeaderFragment;
import org.junit.Before;
import org.junit.Rule;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.mockito.ArgumentCaptor;
import org.mockito.InOrder;
import org.mockito.Mock;
import org.mockito.junit.MockitoJUnitRunner;

import java.util.concurrent.CountDownLatch;

import static com.atlassian.bitbucket.mesh.rpc.v1.RpcSidebandFollowerFragment.SidebandOneofCase.HEARTBEAT;
import static org.junit.Assert.*;
import static org.mockito.Mockito.*;

@RunWith(MockitoJUnitRunner.class)
public class DefaultSidebandManagerTest {

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

    @Mock
    private RequestContext requestContext;
    @Mock
    private RequestManager requestManager;
    @Mock
    private SidebandSender sender;
    private DefaultSidebandManager sidebandManager;

    @Before
    public void setUp() {
        when(requestContext.getAgent()).thenReturn("Control");
        when(requestManager.getContextOrThrow()).thenReturn(requestContext);

        sidebandManager = new DefaultSidebandManager(requestManager, 1);
    }

    @Test
    public void testCloseHandlesAlreadyClosed() {
        doThrow(new IllegalStateException("call is closed")).when(sender).close();

        //The IllegalStateException should be caught and handled. If it propagates, it'll fail the test
        try (SidebandReceiver ignored = sidebandManager.open(openParameters())) {
            assertTrue(sidebandManager.isOpen());
        }

        verify(sender).close();
    }

    @Test
    public void testCloseHandlesOtherExceptions() {
        doThrow(RuntimeException.class).when(sender).close();

        //The RuntimeException should be caught and handled. If it propagates, it'll fail the test
        try (SidebandReceiver ignored = sidebandManager.open(openParameters())) {
            assertTrue(sidebandManager.isOpen());
        }

        verify(sender).close();
    }

    @Test
    public void testCloseTwice() {
        try (SidebandReceiver receiver = sidebandManager.open(openParameters())) {
            //noinspection RedundantExplicitClose because the extra close is what we're testing
            receiver.close();
        }

        verify(sender).close(); //Should only be called once, even though receiver is closed twice
    }

    @Test(timeout = 5000L)
    public void testConcurrentOpenAndSendSendsAllFragments() throws InterruptedException {
        CountDownLatch opening = new CountDownLatch(1);
        CountDownLatch sending = new CountDownLatch(1);

        //The opening thread will send a queued heartbeat (called "first", created/queued below). When
        //that happens, it will countDown the opening latch, which will allow starting the sending
        //thread, and then block on the sending latch
        doAnswer(invocation -> {
            opening.countDown();
            sending.await();

            return null;
        }).when(sender).send(argThat(fragment -> fragment.getSidebandOneofCase() == HEARTBEAT));

        RpcSidebandFollowerFragment first = RpcSidebandFollowerFragment.newBuilder()
                .setHeartbeat(RpcHeartbeat.getDefaultInstance())
                .build();
        assertTrue(sidebandManager.send(first)); //Queued is considered sent
        verifyNoInteractions(sender);

        //Start an opener thread, which will open the sideband connection and try to drain the queue
        Thread open = new Thread(() -> sidebandManager.open(openParameters()),
                getClass().getSimpleName() + "-tco-opener");
        open.setDaemon(true);
        open.start();

        //Wait for the opener thread to open the sideband channel and drain the queue. This way we
        //know that thread has acquired the "channels" monitor
        opening.await();

        //Start a sender thread that wants to send a fragment
        RpcSidebandFollowerFragment second = RpcSidebandFollowerFragment.newBuilder()
                .setRepositoryReplicaObservation(RpcRepositoryReplicaObservation.getDefaultInstance())
                .build();
        Thread send = new Thread(() -> sidebandManager.send(second), getClass().getSimpleName() + "-tco-sender");
        send.setDaemon(true);
        send.start();

        //Wait for the sender thread to block on the "channels" monitor
        assertThreadInState("Sender should be blocked", Thread.State.BLOCKED, send, 1000L);
        //Release the opener thread, which should clear the "channels" monitor and allow the sender
        //thread to send its fragment
        sending.countDown();

        try {
            //Both threads should now complete
            open.join();
            send.join();
        } finally {
            sidebandManager.onGrpcServerShutdown(new GrpcServerShutdownEvent()); //Just to close the sideband
        }

        //Verify both fragments were sent
        InOrder ordered = inOrder(sender);
        ordered.verify(sender).send(same(first));
        ordered.verify(sender).send(same(second));
    }

    @Test
    public void testOnGrpcServerShutdown() {
        sidebandManager.open(openParameters());

        sidebandManager.onGrpcServerShutdown(new GrpcServerShutdownEvent());
        assertFalse(sidebandManager.isOpen());
        verify(sender).close();
    }

    @Test
    public void testOnGrpcServerShutdownWithoutChannels() {
        sidebandManager.onGrpcServerShutdown(new GrpcServerShutdownEvent());
    }

    @Test
    public void testOnReceiveHeartbeat() {
        try (SidebandReceiver receiver = sidebandManager.open(openParameters())) {
            receiver.onReceive(RpcSidebandLeaderFragment.newBuilder()
                    .setHeartbeat(RpcHeartbeat.getDefaultInstance())
                    .build());

            ArgumentCaptor<RpcSidebandFollowerFragment> fragmentCaptor =
                    ArgumentCaptor.forClass(RpcSidebandFollowerFragment.class);
            verify(sender).send(fragmentCaptor.capture());

            RpcSidebandFollowerFragment fragment = fragmentCaptor.getValue();
            assertTrue(fragment.hasHeartbeat());
        }
    }

    @Test
    public void testOpenAndClose() {
        SidebandReceiver receiver = sidebandManager.open(openParameters());
        assertNotNull(receiver);
        assertTrue(sidebandManager.isOpen());
        verifyNoInteractions(sender);

        receiver.close();
        assertFalse(sidebandManager.isOpen());
        verify(sender).close();
    }

    @Test
    public void testOpenWithNullParameters() {
        assertThrows(NullPointerException.class, () -> sidebandManager.open(null));
    }

    @Test
    public void testOpenWithoutRequest() {
        reset(requestManager);
        doThrow(FailedPreconditionException.class).when(requestManager).getContextOrThrow();

        assertThrows(FailedPreconditionException.class, () -> sidebandManager.open(openParameters()));
    }

    @Test
    public void testSend() {
        RpcSidebandFollowerFragment shouldQueue = RpcSidebandFollowerFragment.getDefaultInstance();
        assertTrue(sidebandManager.send(shouldQueue)); //Queued is considered sent

        try (SidebandReceiver ignored = sidebandManager.open(openParameters())) {
            //After the channel is opened, any fragment(s) we tried to send before it was opened
            //should be sent
            verify(sender).send(same(shouldQueue));

            RpcSidebandFollowerFragment fragment = RpcSidebandFollowerFragment.newBuilder()
                    .setHeartbeat(RpcHeartbeat.getDefaultInstance())
                    .build();
            assertTrue(sidebandManager.send(fragment));

            verify(sender).send(same(fragment));
        }
    }

    @Test
    public void testSendAfterClose() {
        try (SidebandReceiver receiver = sidebandManager.open(openParameters())) {
            receiver.close();

            receiver.onReceive(RpcSidebandLeaderFragment.newBuilder()
                    .setHeartbeat(RpcHeartbeat.getDefaultInstance())
                    .build());
        }

        verify(sender, never()).send(any());
    }

    @Test
    public void testSendClosesAfterFailure() {
        doThrow(new IllegalStateException("call is closed")).when(sender).send(any());

        try (SidebandReceiver ignored = sidebandManager.open(openParameters())) {
            //The exception from the sender should not be propagated. This is important for ensuring any
            //racy errors from the sideband don't break other processing
            assertFalse(sidebandManager.send(RpcSidebandFollowerFragment.getDefaultInstance()));

            InOrder ordered = inOrder(sender);
            ordered.verify(sender).send(notNull());
            ordered.verify(sender).close();

            //Even though the receiver isn't closed yet, the send error should have removed it from the
            //queue, so the SidebandManager should have no open channels
            assertFalse(sidebandManager.isOpen());
        }
    }

    @Test
    public void testSendDiscardsFragmentsBeyondQueueLimit() {
        //The queue limit is set to 1, so queuing 2 fragments should see the second one dropped
        RpcSidebandFollowerFragment first = RpcSidebandFollowerFragment.newBuilder().build();
        RpcSidebandFollowerFragment second = RpcSidebandFollowerFragment.newBuilder().build();

        assertTrue(sidebandManager.send(first));
        assertFalse(sidebandManager.send(second));

        try (SidebandReceiver ignored = sidebandManager.open(openParameters())) {
            verify(sender).send(same(first));
            verify(sender, never()).send(same(second));
        }

        logListener.assertWarn("queue is full");
    }

    private static void assertThreadInState(String message, Thread.State state, Thread thread, long timeoutMs) {
        long timeoutTimestamp = System.currentTimeMillis() + timeoutMs;

        while (System.currentTimeMillis() < timeoutTimestamp && thread.getState() != state) {
            try {
                Thread.sleep(2);
            } catch (InterruptedException e) {
                // restore the interrupted flag
                Thread.currentThread().interrupt();
            }
        }

        assertEquals(message, state, thread.getState());
    }

    private OpenSidebandParameters openParameters() {
        return new OpenSidebandParameters.Builder(sender).build();
    }
}