package com.atlassian.bitbucket.context;

import com.atlassian.bitbucket.annotation.DataSet;
import com.atlassian.bitbucket.dbunit.DatabaseConnectionFactory;
import org.dbunit.database.IDatabaseConnection;
import org.dbunit.dataset.CompositeDataSet;
import org.dbunit.dataset.DataSetException;
import org.dbunit.dataset.IDataSet;
import org.dbunit.dataset.xml.FlatDtdDataSet;
import org.dbunit.dataset.xml.FlatXmlDataSetBuilder;
import org.dbunit.operation.DatabaseOperation;
import org.hibernate.SessionFactory;
import org.springframework.beans.factory.InitializingBean;
import org.springframework.context.ApplicationContext;
import org.springframework.core.annotation.AnnotationUtils;
import org.springframework.core.io.Resource;
import org.springframework.core.io.ResourceLoader;
import org.springframework.test.context.TestContext;
import org.springframework.test.context.support.AbstractTestExecutionListener;

import java.io.IOException;
import java.lang.reflect.Method;
import java.util.Arrays;

/**
 * Applies the {@link DataSet} annotation during test execution.
 * <p>
 * Note: The methods on this class do full resource management but do not do <i>any</i> exception handling. Any issues
 * related to creating and inserting/deleting data sets can, and should, fail the associated test.
 */
public class DataSetTestExecutionListener extends AbstractTestExecutionListener implements InitializingBean {

    private final DatabaseConnectionFactory connectionFactory;
    private final Resource dtdResource;
    private final SessionFactory sessionFactory;

    private DataSetBuilder dataSetBuilder;

    @SuppressWarnings("unused") //Someday I'm going to make a library out of this...
    public DataSetTestExecutionListener(DatabaseConnectionFactory connectionFactory, SessionFactory sessionFactory) {
        this(connectionFactory, sessionFactory, null);
    }

    public DataSetTestExecutionListener(DatabaseConnectionFactory connectionFactory, SessionFactory sessionFactory,
                                        Resource dtdResource) {
        this.connectionFactory = connectionFactory;
        this.dtdResource = dtdResource;
        this.sessionFactory = sessionFactory;
    }

    /**
     * Creates the {@link DataSetBuilder} which will be used to transform Spring {@code Resource}s into DbUnit
     * {@code IDataSet} implementations. If a DTD {@code Resource} was supplied during construction, the builder
     * will use that DTD to provide metadata when constructing {@code IDataSet}s. Otherwise, it will attempt to
     * "sense" that metadata from the XML used as input.
     *
     * @throws DataSetException if the DTD resource cannot be transformed into an {@code IDataSet}
     * @throws IOException if a stream cannot be opened from the DTD resource
     */
    @Override
    public void afterPropertiesSet() throws DataSetException, IOException {
        if (dtdResource == null) {
            dataSetBuilder = new NoDtdDataSetBuilder();
        } else {
            dataSetBuilder = new DtdDataSetBuilder(new FlatDtdDataSet(dtdResource.getInputStream()));
        }
    }

    /**
     * After the test method has completed, if the {@link DataSet} annotation is present and has been marked for
     * {@link DataSet#delete() deletion}, deletes the rows defined in the data set(s).
     *
     * @param testContext the test execution context
     * @throws Exception if resources cannot be loaded, a connection cannot be established, rows cannot be deleted or
     *                   any other issue prevents proper execution
     */
    @Override
    public void afterTestMethod(TestContext testContext) throws Exception {
        DataSet annotation = getDataSetAnnotation(testContext);
        if (annotation != null && annotation.delete()) {
            applyAnnotation(testContext, annotation, DatabaseOperation.DELETE, true);
        }

        clearCache();
    }

    /**
     * Before the test method runs, if the {@link DataSet} annotation is present and has been marked for
     * {@link DataSet#insert() insertion}, inserts the rows defined in the data set(s).
     *
     * @param testContext the test execution context
     * @throws Exception if resources cannot be loaded, a connection cannot be established, rows cannot be inserted
     *                   or any other issue prevents proper execution
     */
    @Override
    public void beforeTestMethod(TestContext testContext) throws Exception {
        DataSet annotation = getDataSetAnnotation(testContext);
        if (annotation != null && annotation.insert()) {
            applyAnnotation(testContext, annotation, DatabaseOperation.INSERT, false);
        }
    }

    private void applyAnnotation(TestContext testContext, DataSet annotation, DatabaseOperation operation,
                                 boolean commit) throws Exception {
        ApplicationContext applicationContext = testContext.getApplicationContext();
        String[] locations = computeLocations(testContext, annotation);
        Resource[] resources = resolveResources(applicationContext, locations);
        IDataSet dataSet = createDataSet(resources);

        IDatabaseConnection connection = connectionFactory.newConnection();
        try {
            operation.execute(connection, dataSet);
            if (commit && !connection.getConnection().getAutoCommit()) {
                connection.getConnection().commit();
            }
        } finally {
            connection.close();
        }
    }

    /**
     * Evicts everything from L2 cache, so that tests don't interfere with each other.
     */
    private void clearCache() {
        //Evict all entities from the L2 cache as well
        sessionFactory.getCache().evictEntityRegions();
        //Also evict any collections that have been cached
        sessionFactory.getCache().evictCollectionRegions();
    }

    private String[] computeLocations(TestContext testContext, DataSet dataSet) {
        String[] locations = dataSet.value();
        String[] defaultLocations = (String[]) AnnotationUtils.getDefaultValue(dataSet);
        if (Arrays.equals(locations, defaultLocations)) {
            Class<?> testClass = testContext.getTestClass();

            return new String[]{"classpath:dbunit/" + testClass.getSimpleName() + ".xml"};
        }
        return locations;
    }

    private IDataSet createDataSet(Resource[] resources) throws DataSetException, IOException {
        IDataSet[] dataSets = new IDataSet[resources.length];
        for (int i = 0; i < dataSets.length; ++i) {
            dataSets[i] = dataSetBuilder.build(resources[i]);
        }

        if (dataSets.length == 1) {
            return dataSets[0];
        }
        return new CompositeDataSet(dataSets);
    }

    private DataSet getDataSetAnnotation(TestContext testContext) {
        Method testMethod = testContext.getTestMethod();

        DataSet dataSet = testMethod.getAnnotation(DataSet.class);
        if (dataSet == null) {
            Class<?> testClass = testContext.getTestClass();
            dataSet = testClass.getAnnotation(DataSet.class);
        }
        return dataSet;
    }

    private Resource[] resolveResources(ResourceLoader resourceLoader, String[] locations) {
        Resource[] resources = new Resource[locations.length];
        for (int i = 0; i < locations.length; ++i) {
            resources[i] = resourceLoader.getResource(locations[i]);
        }
        return resources;
    }

    private interface DataSetBuilder {

        IDataSet build(Resource resource) throws DataSetException, IOException;
    }

    private static class DtdDataSetBuilder implements DataSetBuilder {

        private final IDataSet dtd;

        private DtdDataSetBuilder(IDataSet dtd) {
            this.dtd = dtd;
        }

        @Override
        public IDataSet build(Resource resource) throws DataSetException, IOException {
            return new FlatXmlDataSetBuilder()
                    .setMetaDataSet(dtd)
                    .build(resource.getInputStream());
        }
    }

    private static class NoDtdDataSetBuilder implements DataSetBuilder {

        @Override
        public IDataSet build(Resource resource) throws DataSetException, IOException {
            return new FlatXmlDataSetBuilder()
                    .setColumnSensing(true)
                    .setDtdMetadata(false)
                    .build(resource.getInputStream());
        }
    }
}
