package com.atlassian.bitbucket.mesh.coordination;

import com.atlassian.bitbucket.mesh.git.exception.AbortedException;
import com.atlassian.bitbucket.mesh.network.MeshNode;
import com.atlassian.bitbucket.mesh.repository.Repository;
import com.google.common.collect.ImmutableList;
import org.junit.Before;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.mockito.Mock;
import org.mockito.junit.MockitoJUnitRunner;

import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.TimeUnit;

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

@RunWith(MockitoJUnitRunner.class)
public class DefaultWriteLockTest {

    private static final long TIMESTAMP = 123123L;

    @Mock
    private Runnable action;
    @Mock
    private MeshNode node1;
    @Mock
    private MeshNode node2;
    @Mock
    private MeshNode node3;
    @Mock
    private WriteToken.Promise promise;
    @Mock
    private Repository repository;
    private DefaultWriteLock writeLock;

    @Before
    public void setUp() {
        when(node1.getId()).thenReturn("node-1");
        when(node2.getId()).thenReturn("node-2");
        when(node3.getId()).thenReturn("node-3");

        writeLock = new DefaultWriteLock(repository, TIMESTAMP, promise, node1, ImmutableList.of(node2, node3));
    }

    @Test
    public void testAwait() throws InterruptedException {
        CountDownLatch startedLatch = new CountDownLatch(1);
        CountDownLatch lockAcquiredLatch = new CountDownLatch(1);
        Thread awaitingThread = new Thread(() -> {
            try {
                startedLatch.countDown();
                writeLock.await();
                lockAcquiredLatch.countDown();
            } catch (InterruptedException interruptedException) {
                Thread.currentThread().interrupt();
                fail("Interrupted while waiting for write lock");
            }
        });
        awaitingThread.start();
        startedLatch.await();

        //Grant 2 of the 3 tokens
        assertTrue(writeLock.onTokenGranted(node1));
        assertTrue(writeLock.onTokenGranted(node2));

        //Verify that the awaitingThread remains blocked
        assertFalse(lockAcquiredLatch.await(50, TimeUnit.MILLISECONDS));

        assertTrue(writeLock.onTokenGranted(node3));
        assertTrue(lockAcquiredLatch.await(1, TimeUnit.SECONDS));
    }

    @Test
    public void testAwaitThrowsWhenLockIsClosed() throws InterruptedException {
        CountDownLatch startedLatch = new CountDownLatch(1);
        CountDownLatch lockAcquiredLatch = new CountDownLatch(1);
        CountDownLatch abortedLatch = new CountDownLatch(1);
        Thread awaitingThread = new Thread(() -> {
            try {
                startedLatch.countDown();
                writeLock.await();
                lockAcquiredLatch.countDown();
            } catch (AbortedException exception) {
                abortedLatch.countDown();
            } catch (InterruptedException interruptedException) {
                Thread.currentThread().interrupt();
            }
        });
        awaitingThread.start();
        startedLatch.await();

        //Verify that the awaitingThread remains blocked
        assertFalse(lockAcquiredLatch.await(50, TimeUnit.MILLISECONDS));
        assertFalse(abortedLatch.await(0, TimeUnit.MILLISECONDS));

        //Close the lock and verify that await throws an AbortedException
        writeLock.close();
        assertTrue(abortedLatch.await(1, TimeUnit.SECONDS));
    }

    @Test
    public void testCloseBeforeLockIsAcquired() {
        writeLock.close();

        verify(promise).close();

        ExecutionException exception = assertThrows(ExecutionException.class,
                () -> writeLock.asCompletionStage().toCompletableFuture().get());
        assertThat(exception.getCause(), instanceOf(AbortedException.class));
    }

    @Test
    public void testOnTokenGrantIsRejectedAfterClosed() {
        writeLock.close();

        assertFalse(writeLock.onTokenGranted(node1));
    }

    @Test
    public void testOnTokenGrantedFromUnexpectedNodeIsRejected() {
        MeshNode other = mock(MeshNode.class);
        when(other.getId()).thenReturn("other");
        assertFalse(writeLock.onTokenGranted(other));
    }

    @Test
    public void testRemove() {
        assertTrue(writeLock.onTokenGranted(node2));
        assertTrue(writeLock.remove(node2));
    }

    @Test
    public void testRemoveGrantsLockIfAllRemainingTokensAreGranted() {
        assertTrue(writeLock.onTokenGranted(node1));
        assertTrue(writeLock.onTokenGranted(node3));
        assertFalse(writeLock.isAcquired()); //node2's token is still required

        assertFalse(writeLock.remove(node2));
        assertTrue(writeLock.isAcquired()); //node2's token is no longer required, so the lock is now held
    }

    @Test
    public void testRemoveWhenNotHeld() {
        assertFalse(writeLock.remove(node2));

        //Even though remove returned false, that just meant node2's token wasn't held. But it still
        //should have removed node2 as required, which means acquiring tokens for node1 and node3
        //should be enough to acquire the lock
        assertTrue(writeLock.onTokenGranted(node1));
        assertTrue(writeLock.onTokenGranted(node3));
        assertTrue(writeLock.isAcquired());
    }

    @Test
    public void testRemoveWhileLocked() {
        assertTrue(writeLock.onTokenGranted(node1));
        assertTrue(writeLock.onTokenGranted(node2));
        assertTrue(writeLock.onTokenGranted(node3));
        assertTrue(writeLock.isAcquired());

        //Even though we hold the lock, we should still be able to evict a node
        assertTrue(writeLock.remove(node3));
        //And after we do we should still hold the lock
        assertTrue(writeLock.isAcquired());
    }

    @Test
    public void testWriteLockLifecycle() throws ExecutionException, InterruptedException {
        writeLock.asCompletionStage().thenRun(action);

        //0 tokens held
        assertFalse(writeLock.isAcquired());
        verifyNoInteractions(action);

        assertTrue(writeLock.onTokenGranted(node1));

        //1 token held
        assertFalse(writeLock.isAcquired());
        verifyNoInteractions(action);

        assertTrue(writeLock.onTokenGranted(node2));

        //2 tokens held
        assertFalse(writeLock.isAcquired());
        verifyNoInteractions(action);

        assertTrue(writeLock.onTokenGranted(node3));

        //3 tokens held, which acquires the lock and triggers the action
        assertTrue(writeLock.isAcquired());
        verify(action).run();

        writeLock.close();

        //The promise should be closed
        verify(promise).close();
        //The CompletionStage should have completed normally, so getting the result should not throw
        writeLock.asCompletionStage().toCompletableFuture().get();
    }

    @Test
    public void testYieldTokenBeforeLockAcquired() {
        assertTrue(writeLock.onTokenGranted(node1));

        assertTrue(writeLock.yieldToken(node1));

        assertTrue(writeLock.onTokenGranted(node2));
        assertFalse(writeLock.isAcquired());

        assertTrue(writeLock.onTokenGranted(node3));
        assertFalse(writeLock.isAcquired());

        assertTrue(writeLock.onTokenGranted(node1));
        assertTrue(writeLock.isAcquired());
    }

    @Test
    public void testYieldTokenIgnoredAfterLockAcquired() {
        //Acquire the lock
        writeLock.onTokenGranted(node1);
        writeLock.onTokenGranted(node2);
        writeLock.onTokenGranted(node3);

        assertFalse(writeLock.yieldToken(node1));
    }

    @Test
    public void testYieldTokenThatIsNotHeld() {
        assertFalse(writeLock.yieldToken(node1));
    }
}