package com.atlassian.stash.internal.concurrent;

import com.atlassian.beehive.ClusterLock;
import com.atlassian.beehive.ClusterLockService;
import org.junit.Before;
import org.junit.Test;

import javax.annotation.Nonnull;

import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.locks.ReentrantLock;

import static org.hamcrest.Matchers.greaterThan;
import static org.junit.Assert.*;
import static org.mockito.Mockito.doThrow;
import static org.mockito.Mockito.spy;
import static org.mockito.Mockito.when;

public class DefaultMonitoredClusterLockServiceTest {

    private SimpleClusterLockService delegateLockService;
    private DefaultMonitoredClusterLockService lockService;

    @Before
    public void setUp() {
        delegateLockService = spy(new SimpleClusterLockService());
        lockService = new DefaultMonitoredClusterLockService(delegateLockService, true);
    }

    @Test
    public void testNotRecordingIfStatisticsDisabled() {
        lockService = new DefaultMonitoredClusterLockService(new SimpleClusterLockService(), false);

        ClusterLock lock = lockService.getLockForName("lock1");
        lock.lock();

        assertEquals(0, lockService.getLockedCount());
        assertEquals(0, lockService.getTotalAcquiredCount());
        assertEquals(0, lockService.getTotalReleasedCount());

        lock.unlock();
        assertEquals(0, lockService.getLockedCount());
        assertEquals(0, lockService.getTotalAcquiredCount());
        assertEquals(0, lockService.getTotalReleasedCount());
    }

    @Test
    public void testLockAndUnlock() throws InterruptedException {
        ClusterLock lock1 = lockService.getLockForName("lock1");
        ClusterLock lock2 = lockService.getLockForName("lock2");

        lock1.lock();

        assertEquals(1, lockService.getLockedCount());
        assertEquals(1, lockService.getTotalAcquiredCount());
        assertEquals(0, lockService.getTotalReleasedCount());

        lock2.tryLock();
        assertEquals(2, lockService.getLockedCount());
        assertEquals(2, lockService.getTotalAcquiredCount());
        assertEquals(0, lockService.getTotalReleasedCount());

        lock1.unlock();
        assertEquals(1, lockService.getLockedCount());
        assertEquals(2, lockService.getTotalAcquiredCount());
        assertEquals(1, lockService.getTotalReleasedCount());

        lock2.unlock();
        assertEquals(0, lockService.getLockedCount());
        assertEquals(2, lockService.getTotalAcquiredCount());
        assertEquals(2, lockService.getTotalReleasedCount());

        lock1.lockInterruptibly();
        assertEquals(1, lockService.getLockedCount());
        assertEquals(3, lockService.getTotalAcquiredCount());
        assertEquals(2, lockService.getTotalReleasedCount());

        lock1.unlock();
        assertEquals(0, lockService.getLockedCount());
        assertEquals(3, lockService.getTotalAcquiredCount());
        assertEquals(3, lockService.getTotalReleasedCount());
    }

    @Test
    public void testLockErrorCounts() {
        ClusterLock lockErrorLockSpy = spy(new SimpleClusterLock());
        ClusterLock unlockErrorLockSpy = spy(new SimpleClusterLock());
        doThrow(new IllegalArgumentException()).when(lockErrorLockSpy).lock();
        doThrow(new IllegalArgumentException()).when(unlockErrorLockSpy).unlock();

        when(delegateLockService.getLockForName("lockError")).thenReturn(lockErrorLockSpy);
        when(delegateLockService.getLockForName("unlockError")).thenReturn(unlockErrorLockSpy);

        ClusterLock lock = lockService.getLockForName("lock");
        ClusterLock lockErrorLock = lockService.getLockForName("lockError");
        ClusterLock unlockErrorLock = lockService.getLockForName("unlockError");

        // lock 'm up
        lock.lock();
        try {
            lockErrorLock.lock();
        } catch (IllegalArgumentException e) {
            // expected - configured in spy
        }
        unlockErrorLock.lock();

        assertEquals(2, lockService.getLockedCount());
        assertEquals(2, lockService.getTotalAcquiredCount());
        assertEquals(1, lockService.getTotalAcquireErrorCount());
        assertEquals(0, lockService.getTotalReleasedCount());
        assertEquals(0, lockService.getTotalReleaseErrorCount());

        // and release 'm again
        lock.unlock();
        try {
            lockErrorLock.unlock();
        } catch (IllegalMonitorStateException e) {
            // expected - not locked
        }
        try {
            unlockErrorLock.unlock();
        } catch (IllegalArgumentException e) {
            // expected - configured in spy
        }

        assertEquals(1, lockService.getLockedCount());
        assertEquals(2, lockService.getTotalAcquiredCount());
        assertEquals(1, lockService.getTotalAcquireErrorCount());
        assertEquals(1, lockService.getTotalReleasedCount());
        assertEquals(2, lockService.getTotalReleaseErrorCount());
    }

    @Test(timeout = 5000L)
    public void testWaitingForLockCount() throws InterruptedException {
        CountDownLatch startLatch = new CountDownLatch(5);
        CountDownLatch endLatch = new CountDownLatch(5);
        ClusterLock lock = lockService.getLockForName("lock");

        lock.lock();
        assertEquals(0, lockService.getQueuedThreadCount());

        Thread thread1 = new Thread(() -> {
            // should block waiting indefinitely
            startLatch.countDown();
            lock.lock();
            lock.unlock();
            endLatch.countDown();
        });
        Thread thread2 = new Thread(() -> {
            // should not block at all
            startLatch.countDown();
            if (lock.tryLock()) {
                lock.unlock();
            }
            endLatch.countDown();
        });
        Thread thread3 = new Thread(() -> {
            try {
                startLatch.countDown();
                if (lock.tryLock(1, TimeUnit.MINUTES)) {
                    lock.unlock();
                }
                endLatch.countDown();
            } catch (InterruptedException e) {
                fail();
            }
        });
        Thread thread4 = new Thread(() -> {
            try {
                startLatch.countDown();
                if (lock.tryLock(10, TimeUnit.MILLISECONDS)) {
                    lock.unlock();
                }
                endLatch.countDown();
            } catch (InterruptedException e) {
                fail();
            }
        });
        Thread thread5 = new Thread(() -> {
            try {
                startLatch.countDown();
                lock.lockInterruptibly();
                lock.unlock();
                endLatch.countDown();
            } catch (InterruptedException e) {
                fail();
            }
        });

        thread1.start();
        thread2.start();
        thread3.start();
        thread4.start();
        thread5.start();

        startLatch.await();

        // let threads 1 and 3 enter the blocked state
        Thread.sleep(100);

        assertEquals(3, lockService.getQueuedThreadCount());
        assertEquals(1, lockService.getLockedCount());
        assertEquals(1, lockService.getTotalAcquiredCount());
        assertEquals(0, lockService.getTotalAcquireErrorCount());
        assertEquals(0, lockService.getTotalReleasedCount());
        assertEquals(0, lockService.getTotalReleaseErrorCount());

        // unlock and let the 2 blocked threads obtain the lock
        lock.unlock();

        endLatch.await();

        assertEquals(0, lockService.getQueuedThreadCount());
        assertEquals(0, lockService.getLockedCount());
        assertEquals(4, lockService.getTotalAcquiredCount());
        assertEquals(0, lockService.getTotalAcquireErrorCount());
        assertEquals(4, lockService.getTotalReleasedCount());
        assertEquals(0, lockService.getTotalReleaseErrorCount());

        assertThat(lockService.getTotalAcquireTimeMillis(), greaterThan(10L));
    }

    private static class SimpleClusterLockService implements ClusterLockService {

        private final ConcurrentMap<String, ClusterLock> locks = new ConcurrentHashMap<>();

        @Override
        public ClusterLock getLockForName(@Nonnull String lockName) {
            return locks.computeIfAbsent(lockName, __ -> new SimpleClusterLock());
        }
    }

    private static class SimpleClusterLock extends ReentrantLock implements ClusterLock {
    }
}