package com.atlassian.stash.internal.db;

import com.atlassian.bitbucket.i18n.I18nService;
import com.atlassian.bitbucket.i18n.KeyedMessage;
import com.atlassian.bitbucket.i18n.SimpleI18nService;
import com.atlassian.bitbucket.util.Drainable;
import com.atlassian.bitbucket.util.ForcedDrainable;
import com.atlassian.stash.internal.backup.liquibase.LiquibaseDao;
import com.atlassian.stash.internal.hibernate.DelegatingSwappableSessionFactory;
import com.atlassian.stash.internal.hibernate.SwappableSessionFactory;
import com.atlassian.stash.internal.jdbc.DataSourceConfiguration;
import com.atlassian.stash.internal.jdbc.DelegatingSwappableDataSource;
import com.atlassian.stash.internal.jdbc.MutableDataSourceConfiguration;
import com.atlassian.stash.internal.jdbc.SwappableDataSource;
import com.atlassian.stash.internal.maintenance.latch.LatchMode;
import com.atlassian.stash.internal.maintenance.latch.LatchState;
import com.atlassian.stash.internal.migration.MigrationException;
import com.atlassian.stash.internal.migration.MigrationValidationException;
import com.hazelcast.core.Cluster;
import com.hazelcast.core.IExecutorService;
import liquibase.exception.LiquibaseException;
import org.hibernate.engine.spi.SessionFactoryImplementor;
import org.junit.Before;
import org.junit.Rule;
import org.junit.Test;
import org.junit.rules.TemporaryFolder;
import org.junit.runner.RunWith;
import org.mockito.Mock;
import org.mockito.Spy;
import org.mockito.runners.MockitoJUnitRunner;
import org.springframework.beans.factory.BeanCreationException;
import org.springframework.context.ApplicationContext;
import org.springframework.dao.DataRetrievalFailureException;
import org.springframework.jdbc.CannotGetJdbcConnectionException;

import javax.sql.DataSource;
import java.io.Closeable;
import java.io.IOException;
import java.sql.SQLException;
import java.util.concurrent.CyclicBarrier;
import java.util.concurrent.TimeUnit;

import static com.atlassian.stash.internal.db.DefaultDatabaseManager.PROTOTYPE_DATA_SOURCE;
import static com.atlassian.stash.internal.db.DefaultDatabaseManager.PROTOTYPE_SESSION_FACTORY;
import static org.junit.Assert.*;
import static org.mockito.Mockito.*;

@RunWith(MockitoJUnitRunner.class)
public class DefaultDatabaseManagerTest {

    @Rule
    public TemporaryFolder folder = new TemporaryFolder();

    @Mock
    private ApplicationContext applicationContext;
    @Mock
    private Cluster cluster;
    @Mock
    private IExecutorService executorService;
    @Spy
    private I18nService i18nService = new SimpleI18nService();
    @Mock
    private LiquibaseDao liquibaseDao;
    @Mock
    private DatabaseValidator databaseValidator;
    @Mock
    private MutableDataSourceConfiguration configuration;
    @Mock
    private DataSourceConfiguration configurationCopy;
    private DefaultDatabaseManager manager;
    private SwappableDataSource dataSource;
    private SwappableSessionFactory sessionFactory;
    @Mock(extraInterfaces = {Drainable.class, ForcedDrainable.class})
    private DataSource realDataSource;
    @Mock
    private SessionFactoryImplementor realSessionFactory;

    @Before
    public void setUp() {
        when(configuration.copy()).thenReturn(configurationCopy);

        dataSource = spy(new DelegatingSwappableDataSource(realDataSource));
        sessionFactory = spy(new DelegatingSwappableSessionFactory(realSessionFactory));
        manager = spy(new DefaultDatabaseManager(applicationContext, cluster, databaseValidator, configuration,
                executorService, i18nService, liquibaseDao, dataSource, sessionFactory));
    }

    @Test
    public void testAcquireLatch() {
        assertFalse(manager.isLatched());
        DatabaseLatch latch = manager.acquireLatch(LatchMode.LOCAL);

        assertNotNull(latch);
        assertSame(latch, manager.getCurrentLatch());
        assertTrue(manager.isLatched());
    }

    @Test(expected = IllegalStateException.class)
    public void testAcquireLatchWhenLatched() {
        manager.acquireLatch(LatchMode.LOCAL);
        manager.acquireLatch(LatchMode.LOCAL);
    }

    @Test
    public void testDrain() {
        DatabaseLatch latch = manager.acquireLatch(LatchMode.LOCAL);

        // simulate the dataSource not draining on the first attempt, but draining on the second attempt
        when(((Drainable) realDataSource).drain(anyLong(), any(TimeUnit.class))).thenReturn(false, true);

        assertFalse(latch.drain(1, TimeUnit.SECONDS));
        assertTrue(latch.drain(1, TimeUnit.SECONDS));
    }

    @Test
    public void testForceDrain() {
        DatabaseLatch latch = manager.acquireLatch(LatchMode.LOCAL);

        // simulate the dataSource not draining on the first attempt, but draining on the second attempt
        when(((Drainable) realDataSource).drain(anyLong(), any(TimeUnit.class))).thenReturn(false);
        when(((ForcedDrainable) realDataSource).forceDrain(anyLong(), any(TimeUnit.class))).thenReturn(false, true);

        // simulate the dataSource not force draining on the first attempt, but force draining on the second attempt
        assertFalse(latch.forceDrain(1, TimeUnit.SECONDS));
        assertTrue(latch.forceDrain(1, TimeUnit.SECONDS));
    }

    @Test
    public void testDrainWhenUnlatched() {
        DatabaseLatch latch = manager.acquireLatch(LatchMode.LOCAL);
        latch.unlatch();

        // this shouldn't throw an exception. Because of clustering it's entirely possible that one node is draining
        // while another node initiates the unlatch operation.
        latch.drain(0, TimeUnit.SECONDS);
    }

    @Test
    public void testGetHandle() {
        DatabaseHandle handle = manager.getHandle();

        assertSame(configurationCopy, handle.getConfiguration());
        assertSame(realDataSource, handle.getDataSource());
        assertSame(realSessionFactory, handle.getSessionFactory());
    }

    @Test
    public void testGetHandleWhileLatched() throws Exception {
        DatabaseLatch latch = manager.acquireLatch(LatchMode.LOCAL);

        final CyclicBarrier barrier = new CyclicBarrier(2);
        Thread thread = new Thread(() -> {
            try {
                // to signal the main thread we've gotten to this point
                barrier.await();

                // getting the handle should block the thread until the latch is released
                DatabaseHandle handle = manager.getHandle();

                assertSame(configurationCopy, handle.getConfiguration());
                assertSame(realDataSource, handle.getDataSource());
                assertSame(realSessionFactory, handle.getSessionFactory());

                // to signal the main thread that this thread is no longer blocked
                barrier.await();
            } catch (Exception e) {
                fail("Oh noooos");
            }
        });

        thread.start();
        barrier.await(2, TimeUnit.HOURS);

        // at this point 'thread' is about to retrieve the handle, which should cause it to block
        Thread.sleep(50); // to ensure 'thread' actually hits the block
        // verify that 'thread' hasn't gotten the handle
        assertEquals(0, barrier.getNumberWaiting());

        latch.unlatch();
        barrier.await(2, TimeUnit.SECONDS);
        assertEquals(0, barrier.getNumberWaiting());
    }

    @Test
    public void testGetState() {
        assertEquals(LatchState.AVAILABLE, manager.getState());

        // latch the database
        DatabaseLatch latch = manager.acquireLatch(LatchMode.LOCAL);
        // mock that the database initially isn't drained, but on the next call it is
        when(((Drainable) realDataSource).drain(eq(0L), any(TimeUnit.class))).thenReturn(false, true);
        assertEquals(LatchState.LATCHED, manager.getState());
        assertEquals(LatchState.DRAINED, manager.getState());

        latch.unlatch();
        assertEquals(LatchState.AVAILABLE, manager.getState());
    }

    @Test
    public void testUnlatchTo() {
        DatabaseHandle target = mock(DatabaseHandle.class);
        DataSourceConfiguration targetConfiguration = mock(DataSourceConfiguration.class);
        DataSource targetDataSource = mock(DataSource.class);
        SessionFactoryImplementor targetSessionFactory = mock(SessionFactoryImplementor.class);

        when(target.getConfiguration()).thenReturn(targetConfiguration);
        when(target.getDataSource()).thenReturn(targetDataSource);
        when(target.getSessionFactory()).thenReturn(targetSessionFactory);

        DatabaseLatch latch = manager.acquireLatch(LatchMode.LOCAL);
        latch.unlatchTo(target);

        verify(dataSource).swap(same(targetDataSource));
        verify(sessionFactory).swap(same(targetSessionFactory));
        verify(configuration).update(same(targetConfiguration));
    }

    @Test
    public void testCloseDataSourceSwallowsCloseExceptions() throws Exception {
        CloseableDataSource dataSource = mock(CloseableDataSource.class);
        doThrow(IOException.class).when(dataSource).close();

        manager.closeDataSource(dataSource);

        verify(dataSource).close();
    }

    @Test
    public void testCloseDataSourceSwallowsInvocationExceptions() throws Exception {
        DataSourceWithClose dataSource = mock(DataSourceWithClose.class);
        doThrow(SQLException.class).when(dataSource).close();

        manager.closeDataSource(dataSource);

        verify(dataSource).close();
    }

    @Test
    public void testCloseDataSourceSwallowsReflectionExceptions() throws Exception {
        DataSourceWithInaccessibleClose dataSource = mock(DataSourceWithInaccessibleClose.class);

        manager.closeDataSource(dataSource);

        assertFalse(dataSource.isClosed()); //Should not have actually invoked close()
    }

    @Test
    public void testCloseDataSourceWithCloseable() throws Exception {
        CloseableDataSource dataSource = mock(CloseableDataSource.class);

        manager.closeDataSource(dataSource);

        verify(dataSource).close();
    }

    @Test
    public void testCloseDataSourceWithDataSourceWithCloseMethod() throws Exception {
        DataSourceWithClose dataSource = mock(DataSourceWithClose.class);

        manager.closeDataSource(dataSource);

        verify(dataSource).close();
    }

    @Test
    public void testCloseDataSourceWithUncloseableDataSource() {
        DataSource dataSource = mock(DataSource.class);

        manager.closeDataSource(dataSource);
    }

    @Test
    public void testCreateDataSource() {
        DataSourceConfiguration configuration = mock(DataSourceConfiguration.class);
        DataSource dataSource = mock(DataSource.class);

        when(applicationContext.getBean(eq(PROTOTYPE_DATA_SOURCE), same(configuration))).thenReturn(dataSource);

        assertSame(dataSource, manager.createDataSource(configuration));

        verify(applicationContext).getBean(eq(PROTOTYPE_DATA_SOURCE), same(configuration));
    }

    @Test(expected = MigrationException.class)
    public void testCreateDataSourceThrowsOnBeanCreationFailure() {
        DataSourceConfiguration configuration = mock(DataSourceConfiguration.class);

        doThrow(BeanCreationException.class)
                .when(applicationContext).getBean(eq(PROTOTYPE_DATA_SOURCE), same(configuration));

        try {
            manager.createDataSource(configuration);
        } finally {
            verify(applicationContext).getBean(eq(PROTOTYPE_DATA_SOURCE), same(configuration));
        }
    }

    @Test
    public void testCreateSchema() throws Exception {
        DataSource dataSource = mock(DataSource.class);

        manager.createSchema(dataSource);

        verify(liquibaseDao).createSchema(same(dataSource));
    }

    @Test(expected = MigrationException.class)
    public void testCreateSchemaThrowsOnFailure() throws Exception {
        DataSource dataSource = mock(DataSource.class);

        doThrow(LiquibaseException.class).when(liquibaseDao).createSchema(same(dataSource));

        try {
            manager.createSchema(dataSource);
        } finally {
            verify(liquibaseDao).createSchema(same(dataSource));
        }
    }

    @Test
    public void testCreateSessionFactory() {
        DataSource dataSource = mock(DataSource.class);
        DataSourceConfiguration dataSourceConfiguration = mock(DataSourceConfiguration.class);
        SessionFactoryImplementor sessionFactory = mock(SessionFactoryImplementor.class);

        when(applicationContext.getBean(eq(PROTOTYPE_SESSION_FACTORY), same(dataSource))).thenReturn(sessionFactory);

        manager.createSessionFactory(dataSource, dataSourceConfiguration);

        verify(applicationContext).getBean(eq(PROTOTYPE_SESSION_FACTORY), same(dataSource), same(dataSourceConfiguration));
    }

    @Test(expected = MigrationException.class)
    public void testCreateSessionFactoryThrowsOnBeanCreationFailure() {
        DataSource dataSource = mock(DataSource.class);
        DataSourceConfiguration dataSourceConfiguration = mock(DataSourceConfiguration.class);

        doThrow(BeanCreationException.class)
                .when(applicationContext).getBean(eq(PROTOTYPE_SESSION_FACTORY), same(dataSource), same(dataSourceConfiguration));

        try {
            manager.createSessionFactory(dataSource, dataSourceConfiguration);
        } finally {
            verify(applicationContext).getBean(eq(PROTOTYPE_SESSION_FACTORY), same(dataSource), same(dataSourceConfiguration));
        }
    }

    @Test
    public void testValidateConfiguration() {
        manager.setConnectTimeout(0L);
        manager.validateConfiguration(mockH2Configuration());

        verify(databaseValidator).validate(isA(DataSource.class));
    }

    @Test(expected = MigrationException.class)
    public void testValidateConfigurationThrowsOnConnectionException() {
        doThrow(CannotGetJdbcConnectionException.class).when(databaseValidator).validate(isA(DataSource.class));

        try {
            manager.validateConfiguration(mockH2Configuration());
        } catch (MigrationException e) {
            assertEquals("bitbucket.migration.test.connectfailed", e.getMessageKey());

            throw e;
        } finally {
            verify(databaseValidator).validate(isA(DataSource.class));
        }
    }

    @Test(expected = MigrationException.class)
    public void testValidateConfigurationThrowsOnDataException() {
        doThrow(DataRetrievalFailureException.class).when(databaseValidator).validate(isA(DataSource.class));

        try {
            manager.validateConfiguration(mockH2Configuration());
        } catch (MigrationException e) {
            assertEquals("bitbucket.migration.test.supportunverified", e.getMessageKey());

            throw e;
        } finally {
            verify(databaseValidator).validate(isA(DataSource.class));
        }
    }

    @Test(expected = MigrationValidationException.class)
    public void testValidateConfigurationThrowsOnTargetException() {
        Exception exception = new DatabaseValidationException(new KeyedMessage("", "", ""));
        doThrow(exception).when(databaseValidator).validate(isA(DataSource.class));

        try {
            manager.validateConfiguration(mockH2Configuration());
        } finally {
            verify(databaseValidator).validate(isA(DataSource.class));
        }
    }

    @SuppressWarnings("ConstantConditions")
    @Test(expected = NullPointerException.class)
    public void testValidateConfigurationThrowsOnNullConfiguration() {
        manager.validateConfiguration(null);
    }

    private DataSourceConfiguration mockH2Configuration() {
        DataSourceConfiguration configuration = mock(DataSourceConfiguration.class);
        when(configuration.getDriverClassName()).thenReturn("org.h2.Driver");
        when(configuration.getUrl()).thenReturn("jdbc:h2:file:" +
                folder.getRoot().getAbsolutePath() + "/db;DB_CLOSE_ON_EXIT=TRUE");
        when(configuration.getUser()).thenReturn("sa");
        return configuration;
    }

    private interface CloseableDataSource extends Closeable, DataSource {

        //No methods necessary; this interface just combines two others for testing
    }

    private abstract static class DataSourceWithClose implements DataSource {

        public void close() {
            //No implementation necessary; this will be verified with mocks
        }
    }

    private abstract static class DataSourceWithInaccessibleClose implements DataSource {

        private boolean closed;

        public boolean isClosed() {
            return closed;
        }

        @SuppressWarnings("unused")
        private void close() {
            //This method needs to exist so it can be found via ReflectionUtils.findMethod, but be inaccessible
            //for invocation via Reflection
            closed = true;
        }
    }
}
