/*
 * Copyright 2005-2010 the original author or authors.
 * 
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 * 
 *      http://www.apache.org/licenses/LICENSE-2.0
 * 
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.wamblee.test.persistence;

import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.util.ArrayList;
import java.util.List;
import java.util.logging.Level;
import java.util.logging.Logger;

import javax.sql.DataSource;

import junit.framework.TestCase;

import org.dbunit.DataSourceDatabaseTester;
import org.dbunit.IDatabaseTester;
import org.dbunit.database.DatabaseConnection;
import org.dbunit.database.DatabaseSequenceFilter;
import org.dbunit.database.IDatabaseConnection;
import org.dbunit.dataset.FilteredDataSet;
import org.dbunit.dataset.IDataSet;
import org.dbunit.dataset.filter.ITableFilter;
import org.dbunit.dataset.filter.ITableFilterSimple;
import org.dbunit.operation.DatabaseOperation;

/**
 * Database utilities is a simple support class for common tasks in working with
 * databases.
 */
public class DatabaseUtils {

    /**
     * Represents a set of tables.
     *  
     * @author Erik Brakkee
     */
    public static interface TableSet {
        boolean contains(String aTableName);
    }

    /**
     * Represents a unit of work (transaction).
     * 
     * @author Erik Brakkee
     *
     * @param <T> Type of return value. 
     */
    public static interface JdbcUnitOfWork<T> {
        /**
         * Executes statement within a transaction. 
         * @param aConnection Connection. 
         * @return Result of the work. 
         * @throws Exception
         */
        T execute(Connection aConnection) throws Exception;
    }

    /**
     * Operation to be executed on a set of tables for each table 
     * individually. 
     * 
     * @author Erik Brakkee
     */
    public static interface TableSetOperation {
        /**
         * Executes on a table. 
         * @param aTable Table name. 
         * @throws Exception
         */
        void execute(String aTable) throws Exception;
    }

    private static final Logger LOG = Logger.getLogger(DatabaseUtils.class
        .getName());

    /**
     * Schema pattern.
     */
    private static final String SCHEMA_PATTERN = "%";
    private DataSource dataSource;

    private IDatabaseTester dbtester;
    /**
     * List of connections that were created for dbtesters. 
     * This list will be closed in the {@link #stop()} method. 
     */
    private List<IDatabaseConnection> connections; 
    
    /**
     * Constructs the database utils. 
     * Before use, {@link #start()} must be called. 
     * @param aDataSource Datasource.  
     */
    public DatabaseUtils(DataSource aDataSource) {
        dataSource = aDataSource;
        dbtester = new DataSourceDatabaseTester(dataSource);
        connections = new ArrayList<IDatabaseConnection>(); 
    }

    /**
     * Starts the database utils. 
     */
    public void start() { 
        // Empty. No operation currently. 
    }

    /**
     * Stops the database utils, closing any JDBC connections that were created
     * by this utility. Note that connections obtained from the datasource directly 
     * must still be closed by the user. The involved connections are only those that
     * are created by this utility. 
     */
    public void stop() {
        for (IDatabaseConnection connection: connections) { 
            try {
                connection.close();
            } catch (SQLException e) {
                LOG.log(Level.WARNING, "Could not close connection", e);
            }
        }
        connections.clear();
    }

    /**
     * Creates database tester. 
     * @param aTables Tables to create the tester for. 
     * @return Database tester. 
     * @throws Exception
     */
    public IDatabaseTester createDbTester(ITableFilterSimple aTables) throws Exception {
        return createDbTester(getTableNames(aTables));
    }
    
    /**
     * Creates database tester. 
     * @param aTables Tables to create the tester for. 
     * @return Database tester. 
     * @throws Exception
     */
    public IDatabaseTester createDbTester(String[] aTables) throws Exception {
        IDatabaseConnection connection = dbtester.getConnection();
        connections.add(connection);
        dbtester.setDataSet(connection.createDataSet(aTables));
        return dbtester;
    }

    /**
     * Executes an operation on a set of tables. 
     * @param aTables Tables. 
     * @param aOperation Operation. 
     * @throws Exception
     */
    public void executeOnTables(ITableFilterSimple aTables,
        final TableSetOperation aOperation) throws Exception {
        final String[] tableNames = getTableNames(aTables);
        executeInTransaction(new JdbcUnitOfWork<Void>() {
            public Void execute(Connection aConnection) throws Exception {
                for (int i = tableNames.length - 1; i >= 0; i--) {
                    aOperation.execute(tableNames[i]);
                }
                return null;
            }
        });
    }

    /**
     * Cleans a number of database tables. This means deleting the content not dropping the tables. 
     * This may fail in case of cyclic dependencies between the tables (current limitation). 
     * @param aSelection Tables.
     * @throws Exception
     */
    public void cleanDatabase(ITableFilterSimple aSelection) throws Exception {

        final String[] tableNames = getTableNames(aSelection);
        executeInTransaction(new JdbcUnitOfWork<Void>() {

            public Void execute(Connection aConnection) throws Exception {
                IDatabaseConnection connection = new DatabaseConnection(
                    aConnection);
                ITableFilter filter = new DatabaseSequenceFilter(connection,
                    tableNames);
                IDataSet dataset = new FilteredDataSet(filter, connection
                    .createDataSet(tableNames));
                DatabaseOperation.DELETE_ALL.execute(connection, dataset);
                return null;
            }
        });

    }

    /**
     * Executes a unit of work within a transaction. 
     * @param <T> Result type of th ework. 
     * @param aWork Unit of work. 
     * @return
     * @throws Exception
     */
    public <T> T executeInTransaction(JdbcUnitOfWork<T> aWork)
        throws Exception {
        Connection connection = dataSource.getConnection();
        connection.setAutoCommit(false);
        try {
            T value = aWork.execute(connection);
            connection.commit();
            return value;
        } finally {
            connection.close();
        }
    }
    
    /**
     * Returns table names based on a table filter. 
     * @param aSelection Table filter. 
     * @return Table names. 
     * @throws Exception
     */
    public String[] getTableNames(ITableFilterSimple aSelection)
        throws Exception {

        List<String> result = new ArrayList<String>();
        LOG.fine("Getting database table names to clean (schema: '" +
            SCHEMA_PATTERN + "'");

        Connection connection = dataSource.getConnection();
        try {
            ResultSet tableNames = connection.getMetaData().getTables(null,
                SCHEMA_PATTERN, "%", new String[] { "TABLE" });
            while (tableNames.next()) {
                String table = tableNames.getString("TABLE_NAME");
                if (aSelection.accept(table)) {
                    result.add(table);
                }
            }
            return (String[]) result.toArray(new String[0]);
        } finally {
            connection.close();
        }
    }

    /**
     * Use {@link #cleanDatabase(ITableFilterSimple)} instead. 
     */
    @Deprecated
    public void emptyTables(final ITableFilterSimple aSelection)
        throws Exception {
        executeOnTables(aSelection, new TableSetOperation() {
            public void execute(String aTable) throws Exception {
                emptyTable(aTable);
            }
        });
    }

    /**
     * User {@link #cleanDatabase(ITableFilterSimple)} instead. 
     */
    @Deprecated
    public void emptyTable(String aTable) throws Exception {
        executeSql("delete from " + aTable);
    }

    /**
     * Drops tables. This only works if there are no cyclic dependencies between the tables. 
     * @param aTables Tables to drop. 
     * @throws Exception
     */
    public void dropTables(ITableFilterSimple aTables) throws Exception {
        final String[] tableNames = getTableNames(aTables);
        String[] sortedTables = executeInTransaction(new JdbcUnitOfWork<String[]>() {

            public String[] execute(Connection aConnection) throws Exception {
                IDatabaseConnection connection = new DatabaseConnection(
                    aConnection);
                ITableFilter filter = new DatabaseSequenceFilter(connection,
                    tableNames);
                IDataSet dataset = new FilteredDataSet(filter, connection
                    .createDataSet(tableNames));
                return dataset.getTableNames();
            }
        });
        for (int i = sortedTables.length - 1; i >= 0; i--) {
            dropTable(sortedTables[i]);
        }
    }

    /**
     * Drops a table. 
     * @param aTable Table to drop. 
     * @throws Exception
     */
    public void dropTable(final String aTable) throws Exception {
        executeInTransaction(new JdbcUnitOfWork<Void>() {
            public Void execute(Connection aConnection) throws Exception {
                executeUpdate(aConnection, "drop table " + aTable);
                return null;
            }
        });

    }

    /**
     * Executes an SQL statement within a transaction.
     * 
     * @param aSql
     *            SQL statement.
     * @return Return code of the corresponding JDBC call.
     */
    public int executeSql(final String aSql) throws Exception {
        return executeSql(aSql, new Object[0]);
    }

    /**
     * Executes an SQL statement within a transaction. See
     * {@link #setPreparedParam(int, PreparedStatement, Object)}for details on
     * supported argument types.
     * 
     * @param aSql
     *            SQL statement.
     * @param aArg
     *            Argument of the sql statement.
     * @return Return code of the corresponding JDBC call.
     */
    public int executeSql(final String aSql, final Object aArg)
        throws Exception {
        return executeSql(aSql, new Object[] { aArg });
    }

    /**
     * Executes an sql statement. See
     * {@link #setPreparedParam(int, PreparedStatement, Object)}for details on
     * supported argument types.
     * 
     * @param aSql
     *            SQL query to execute.
     * @param aArgs
     *            Arguments.
     * @return Number of rows updated.
     */
    public int executeSql(final String aSql, final Object[] aArgs)
        throws Exception {
        return executeInTransaction(new JdbcUnitOfWork<Integer>() {
            public Integer execute(Connection aConnection) throws Exception {
                PreparedStatement stmt = aConnection.prepareStatement(aSql);
                setPreparedParams(aArgs, stmt);
                return stmt.executeUpdate();
            }
        });
    }

    /**
     * Executes an SQL query.
     * 
     * @param aSql
     *            Query to execute.
     * @return Result set.
     */
    public ResultSet executeQuery(Connection aConnection, String aSql) {
        return executeQuery(aConnection, aSql, new Object[0]);
    }

    /**
     * Executes a query with a single argument. See
     * {@link #setPreparedParam(int, PreparedStatement, Object)}for details on
     * supported argument types.
     * 
     * @param aSql
     *            Query.
     * @param aArg
     *            Argument.
     * @return Result set.
     */
    public ResultSet executeQuery(Connection aConnection, String aSql,
        Object aArg) {
        return executeQuery(aConnection, aSql, new Object[] { aArg });
    }

    /**
     * Executes a query. See
     * {@link #setPreparedParam(int, PreparedStatement, Object)}for details on
     * supported argument types.
     * 
     * @param aSql
     *            Sql query.
     * @param aArgs
     *            Arguments to the query.
     * @return Result set.
     */
    public ResultSet executeQuery(Connection aConnection, final String aSql,
        final Object[] aArgs) {
        try {
            PreparedStatement statement = aConnection.prepareStatement(aSql);
            setPreparedParams(aArgs, statement);

            return statement.executeQuery();
        } catch (SQLException e) {
            throw new RuntimeException(e);
        }
    }

    /**
     * Executes an update. 
     * @param aConnection Connection to use. 
     * @param aSql SQL update to use. 
     * @param aArgs Arguments to the update. 
     * @return Number of rows updated. 
     */
    public int executeUpdate(Connection aConnection, final String aSql,
        final Object... aArgs) {
        try {
            PreparedStatement statement = aConnection.prepareStatement(aSql);
            setPreparedParams(aArgs, statement);

            return statement.executeUpdate();
        } catch (SQLException e) {
            throw new RuntimeException(e);
        }
    }

    /**
     * Sets the values of a prepared statement. See
     * {@link #setPreparedParam(int, PreparedStatement, Object)}for details on
     * supported argument types.
     * 
     * @param aArgs
     *            Arguments to the prepared statement.
     * @param aStatement
     *            Prepared statement
     * @throws SQLException
     */
    private void setPreparedParams(final Object[] aArgs,
        PreparedStatement aStatement) throws SQLException {
        for (int i = 1; i <= aArgs.length; i++) {
            setPreparedParam(i, aStatement, aArgs[i - 1]);
        }
    }

    /**
     * Sets a prepared statement parameter.
     * 
     * @param aIndex
     *            Index of the parameter.
     * @param aStatement
     *            Prepared statement.
     * @param aObject
     *            Value Must be of type Integer, Long, or String.
     * @throws SQLException
     */
    private void setPreparedParam(int aIndex, PreparedStatement aStatement,
        Object aObject) throws SQLException {
        if (aObject instanceof Integer) {
            aStatement.setInt(aIndex, ((Integer) aObject).intValue());
        } else if (aObject instanceof Long) {
            aStatement.setLong(aIndex, ((Long) aObject).longValue());
        } else if (aObject instanceof String) {
            aStatement.setString(aIndex, (String) aObject);
        } else {
            TestCase.fail("Unsupported object type for prepared statement: " +
                aObject.getClass() + " value: " + aObject + " statement: " +
                aStatement);
        }
    }

    /**
     * Gets the table size. 
     * @param aTable Table. 
     * @return Table size. 
     * @throws SQLException
     */
    public int getTableSize(final String aTable) throws Exception {
        return executeInTransaction(new JdbcUnitOfWork<Integer>() {
            public Integer execute(Connection aConnection) throws Exception {
                ResultSet resultSet = executeQuery(aConnection,
                    "select count(*) from " + aTable);
                resultSet.next();
                return resultSet.getInt(1);
            }
        });

    }

    /**
     * Counts the results in a result set. 
     * @param aResultSet Resultset. 
     * @return Number of rows in the set. 
     * @throws SQLException
     */
    public int countResultSet(ResultSet aResultSet) throws SQLException {
        int count = 0;

        while (aResultSet.next()) {
            count++;
        }

        return count;
    }

}