package com.atlassian.stash.internal.hikari;

import com.atlassian.bitbucket.util.UncheckedOperation;
import com.atlassian.stash.internal.jdbc.DataSourceConfiguration;
import com.atlassian.stash.internal.jdbc.ConnectionTracker;
import com.atlassian.stash.internal.jdbc.SimpleDataSourceConfiguration;
import org.junit.Before;
import org.junit.Test;
import org.mockito.invocation.InvocationOnMock;
import org.mockito.stubbing.Answer;

import java.sql.Connection;
import java.sql.SQLException;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;
import java.util.function.BiConsumer;

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

public class ExtendedHikariDataSourceTest {

    private ConnectionTracker connectionTracker;
    private ExtendedHikariDataSource dataSource;

    @Before
    public void setup() {
        //We don't want to even try to create real database connections, so if something does, the fake URL
        //used here should make the test fail. Note that the JDBC driver class _must_ be real; Hikari checks
        //it when HikariConfig.setDriverClassName(String) is called
        DataSourceConfiguration configuration = new SimpleDataSourceConfiguration("org.h2.Driver",
                "jdbc:fake://localhost:7990/test", "sa", null);

        connectionTracker = mock(ConnectionTracker.class);

        dataSource = new ExtendedHikariDataSource(configuration);
        dataSource.setConnectionTracker(connectionTracker);
        dataSource.setDrainPollInterval(25); //Milliseconds
    }

    @Test
    public void testClose() {
        dataSource.close();
    }

    @Test
    public void testDrain() throws Exception {
        doAnswer(new Answer<Object>() {

            private int leased = 5;

            @Override
            public Object answer(InvocationOnMock invocation) throws Throwable {
                return leased--;
            }
        }).when(connectionTracker).getConnectionCount();

        //Should loop 6 times, 25ms between, and eventually drain
        assertTrue(dataSource.drain(2, TimeUnit.SECONDS));
        verify(connectionTracker, times(6)).getConnectionCount();
    }

    @Test
    public void testDrainZeroTimeoutWhenNoneLeased() {
        doReturn(0).when(connectionTracker).getConnectionCount();

        assertTrue(dataSource.drain(0, TimeUnit.SECONDS));
        verify(connectionTracker).getConnectionCount(); // verify it's called only once
    }

    @Test
    public void testDrainZeroTimeoutWhenSomeLeased() {
        doReturn(1).when(connectionTracker).getConnectionCount();

        assertFalse(dataSource.drain(0, TimeUnit.SECONDS));
        verify(connectionTracker, atLeastOnce()).getConnectionCount(); // verify it's called at least once
    }

    @Test
    public void testDrainInterrupted() throws Exception {
        testInterrupted(() -> dataSource.drain(10, TimeUnit.SECONDS));
    }

    @Test
    public void testDrainTimingOut() {
        //Set a longer polling interval
        dataSource.setDrainPollInterval(TimeUnit.SECONDS.toMillis(5));

        doReturn(10).when(connectionTracker).getConnectionCount();

        assertFalse(dataSource.drain(50, TimeUnit.MILLISECONDS));
    }

    @Test
    public void testForceDrainWhenNoneLeased() throws Exception {
        doReturn(0).when(connectionTracker).getConnectionCount();
        assertTrue("Force draining the dataSource should succeed when there are no leases", dataSource.forceDrain(0, TimeUnit.SECONDS));
    }

    @Test
    public void testForceDrainInterruptsOwningThreads() throws Exception {
        Connection connection = mock(Connection.class);
        Thread owningThread = spy(new Thread("testForceDrainInterruptsOwningThreads"));
        trainForEach(connection, owningThread);

        //The fourth time we check for leased connections (that is, the first time after interruption, we report none)
        doReturn(1).doReturn(1).doReturn(1).doReturn(0).when(connectionTracker).getConnectionCount();

        assertTrue("Force draining the dataSource should succeed when there are no leases", dataSource.forceDrain(0, TimeUnit.SECONDS));
        verify(owningThread).interrupt();
    }

    @Test
    public void testForceDrainClosesConnections() throws Exception {
        Connection connection = mock(Connection.class);
        trainForEach(connection);

        //The second time we check for leased connections (that is, the first time after interruption, we report none)
        doReturn(1).when(connectionTracker).getConnectionCount();

        assertFalse("Force draining the dataSource should fail when there are still leases", dataSource.forceDrain(0, TimeUnit.SECONDS));
        verify(connection).close();
    }

    @Test
    public void testForceDrainIgnoresExceptions() throws Exception {
        Connection connection = mock(Connection.class);
        doThrow(SQLException.class).when(connection).close();
        trainForEach(connection);

        //The second time we check for leased connections (that is, the first time after interruption, we report none)
        doReturn(1).when(connectionTracker).getConnectionCount();

        assertFalse("Force draining the dataSource should fail when there are still leases", dataSource.forceDrain(0, TimeUnit.SECONDS));
        verify(connection).close();
    }

    @Test
    public void testForceDrainTimingOut() {
        //Set a longer polling interval
        dataSource.setDrainPollInterval(TimeUnit.SECONDS.toMillis(5));

        doReturn(10).when(connectionTracker).getConnectionCount();

        assertFalse(dataSource.forceDrain(50, TimeUnit.MILLISECONDS));
    }

    @Test
    public void testForceDrainInterrupted() throws Exception {
        testInterrupted(() -> dataSource.forceDrain(10, TimeUnit.SECONDS));
    }

    private void testInterrupted(final UncheckedOperation<Boolean> callable) throws Exception {
        //final interrupted status from test thread
        boolean[] interrupted = {false};

        CountDownLatch readyLatch = new CountDownLatch(1);
        CountDownLatch doneLatch = new CountDownLatch(1);

        Thread thread = new Thread(() -> {
            //Notify the main thread that the testing thread has started. This eliminates a race condition where
            //a Thread.sleep() in the main thread may elapse while the testing thread is still starting
            readyLatch.countDown();

            try {
                assertFalse("Draining the dataSource should fail", callable.perform());

                interrupted[0] = Thread.interrupted();
            } finally {
                doneLatch.countDown();
            }
        });

        dataSource.setDrainPollInterval(TimeUnit.SECONDS.toMillis(5));
        doReturn(10).when(connectionTracker).getConnectionCount();
        thread.start();

        //Wait for the testing thread to start
        readyLatch.await();

        //Once the testing thread has start, give it a quarter second to get blocked in its own Thread.sleep()
        Thread.sleep(250);

        //Interrupt the testing thread to trigger an exception
        thread.interrupt();

        doneLatch.await();
        //This assert ensures that the testing thread received an InterruptedException. If it was not asleep when
        //thread.interrupt() was called, this flag won't be set and the test will fail
        assertTrue("Thread should have been interrupted", interrupted[0]);
    }

    private void trainForEach(Connection connection) {
        trainForEach(connection, spy(new Thread(getClass().getName())));
    }

    @SuppressWarnings("unchecked")
    private void trainForEach(Connection connection, Thread thread) {
        doAnswer(invocation -> {
            BiConsumer<Connection, Thread> consumer = (BiConsumer) invocation.getArguments()[0];
            consumer.accept(connection, thread);

            return null;
        }).when(connectionTracker).forEach(isA(BiConsumer.class));
    }
}
