package com.atlassian.stash.internal.backup.liquibase;

import com.atlassian.bitbucket.util.CancelState;
import com.atlassian.fugue.Effect;
import com.atlassian.fugue.Option;
import com.atlassian.stash.internal.liquibase.LiquibaseUtils;
import com.google.common.annotations.VisibleForTesting;
import com.google.common.base.Preconditions;
import com.google.common.base.Supplier;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.Maps;
import liquibase.change.Change;
import liquibase.change.ColumnConfig;
import liquibase.change.core.DeleteDataChange;
import liquibase.change.core.InsertDataChange;
import liquibase.changelog.DatabaseChangeLog;
import liquibase.database.Database;
import liquibase.database.core.OracleDatabase;
import liquibase.database.structure.Column;
import liquibase.database.structure.Table;
import liquibase.exception.DatabaseException;
import liquibase.exception.LiquibaseException;
import liquibase.exception.LockException;
import liquibase.lockservice.LockService;
import liquibase.sql.visitor.SqlVisitor;
import liquibase.statement.DatabaseFunction;
import liquibase.util.JdbcUtils;
import org.apache.commons.lang3.StringUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.context.ApplicationContext;
import org.springframework.context.ApplicationContextAware;
import org.springframework.dao.CannotAcquireLockException;
import org.springframework.dao.DataRetrievalFailureException;
import org.springframework.stereotype.Component;

import javax.annotation.Nonnull;
import javax.sql.DataSource;
import java.math.BigDecimal;
import java.sql.*;
import java.util.*;
import java.util.function.Predicate;
import java.util.stream.Collectors;

import static com.atlassian.bitbucket.util.Chainable.chain;
import static com.atlassian.fugue.Option.none;
import static com.atlassian.fugue.Option.some;
import static com.atlassian.stash.internal.backup.liquibase.DefaultLiquibaseSession.TO_LOWERCASE_COLUMN_NAME;
import static com.atlassian.stash.internal.backup.liquibase.DefaultLiquibaseSession.TO_LOWERCASE_TABLE_NAME;
import static com.google.common.collect.Iterables.find;

@Component("liquibaseDao")
public class DefaultLiquibaseDao implements ApplicationContextAware, LiquibaseDao {

    public static final DatabaseChangeLog EMPTY_CHANGE_LOG = new DatabaseChangeLog();
    public static final List<SqlVisitor> NO_VISITORS = ImmutableList.of();

    private static final Logger log = LoggerFactory.getLogger(DefaultLiquibaseDao.class);

    private static final Predicate<String> IS_HIBERNATE_TEMPORARY_TABLE = tableName -> {
        // Hibernate temp tables should start with ht_, but a bug (HHH-9290) in
        // Oracle8iDialect.generateTemporaryTableName  truncates the 'h' when
        // the table name is longer than 30 characters
        return tableName.startsWith("ht_") ||
                (tableName.length() == 29 && tableName.startsWith("t_"));
    };

    /**
     * A function that releases the lock on a Liquibase service passed to it.
     * Exceptions are logged as errors, but otherwise silently succeed.
     */
    private static final Effect<LockService> RELEASE_LOCK = new Effect<LockService>() {
        @Override
        public void apply(LockService lockService) {
            try {
                lockService.releaseLock();
            } catch (LockException e) {
                log.error("Failed to release Liquibase lock", e);
            }
        }
    };

    /**
     * Encapsulating a connection to the underlying database.
     * <p>
     * It will continue to use the same database connection until its {@link ThreadLocal#remove} method is called. The
     * next {@link ThreadLocal#get} or {@link LiquibaseSession} method call following a call to {@code remove()} will
     * obtain a new database connection.
     */
    private final LiquibaseSessionThreadLocal databaseSession;
    /**
     * The number of inserts within a change set after which a commit will performed.
     * If the value of commitBlockSize is less than one, then changes will not be committed until the end of the
     * change set.
     */
    private final long commitBlockSize;
    /**
     * A None value indicates that the underlying Liquibase instance is not locked. A Some(x) value is used in the
     * {@link #unlock()} method to unlock the underlying Liquibase instance.
     */
    private Option<LockService> lockService = none();

    private ApplicationContext applicationContext;

    /**
     * A number of records to get in memory in one go during backup.
     * Note: JDBC setFetchSize() is a hint to the driver and behaviour is
     * implementation specific. Recent PostgreSQL, MySQL and Oracle drivers support it natively.
     * SQL Server reportedly doesn't support fetch size, but jTDS driver for SQL Server does
     */
    private static final int fetchSize = 1000;

    @Autowired
    public DefaultLiquibaseDao(@Qualifier("backupDataSourceSupplier") Supplier<DataSource> dataSourceSupplier,
                               @Value("${liquibase.commit.block.size}") long commitBlockSize) {
        this.commitBlockSize = commitBlockSize;
        this.databaseSession = new LiquibaseSessionThreadLocal(dataSourceSupplier);
    }

    @Override
    public void beginChangeSet() {
        if (databaseSession.getDatabase().supportsDDLInTransaction()) {
            try {
                // use transactions
                databaseSession.getDatabase().setAutoCommit(false);
            } catch (DatabaseException e) {
                throw new LiquibaseDataAccessException("Failed to set auto-commit off", databaseSession.getDatabase(), e);
            }
        }
    }

    @Override
    public void close() {
        unlock();
        // close the session to release any resources
        databaseSession.close();
        // ensure that we get a new database connection when we next call databaseSession.get().
        databaseSession.remove();
    }

    @Override
    public long countRows(@Nonnull String tableName) {
        Preconditions.checkNotNull(tableName, "tableName");
        String escapedTableName = databaseSession.getDatabase().escapeTableName(databaseSession.getSnapshot().getSchema(), tableName);
        Statement statement = null;
        ResultSet resultSet = null;
        try {
            Connection connection = LiquibaseUtils.getConnection(databaseSession.getDatabase());
            statement = connection.createStatement();
            resultSet = statement.executeQuery("SELECT count(1) FROM " + escapedTableName);
            resultSet.next();
            return resultSet.getLong(1);
        } catch (SQLException e) {
            throw new DataRetrievalFailureException("Could not count rows for " + tableName, e);
        } finally {
            JdbcUtils.closeResultSet(resultSet);
            JdbcUtils.closeStatement(statement);
        }
    }

    @Override
    public void createSchema(DataSource dataSource) throws LiquibaseException {
        //See the comments on the prototype bean in stash-context.xml. This relies on configuration applied in the
        //XML file, ensuring this bean is identically configured to the master "liquibase" bean that manages the
        //schema during startup. Only the DataSource used changes
        applicationContext.getBean("liquibasePrototype", dataSource);
    }

    /**
     * This implementation always commits the changes to the underlying database.
     */
    @Override
    public void deleteAllRows(@Nonnull String tableName) {
        DeleteDataChange change = new DeleteDataChange();
        change.setTableName(tableName);
        applyChange(change);
        commit();
    }

    @Override
    public void endChangeSet() {
        commit();
    }

    @Override
    public Set<Class<?>> findCustomChanges() {
        return LiquibaseUtils.findCustomChanges();
    }

    @Override
    public long forEachRow(@Nonnull String tableName, String orderingColumn, @Nonnull CancelState cancelState,
                           @Nonnull Effect<Map<String, Object>> effect) {
        Preconditions.checkNotNull(tableName, "tableName");
        Preconditions.checkArgument(orderingColumn == null || StringUtils.isNotBlank(orderingColumn), "blank ordering column");
        Preconditions.checkNotNull(effect, "effect");

        Table table = databaseSession.getSnapshot().getTable(tableName);
        String escapedTableName = databaseSession.getDatabase().escapeTableName(databaseSession.getSnapshot().getSchema(), tableName);
        long numberOfRows = 0;
        Statement statement = null;
        ResultSet resultSet = null;
        try {
            Connection connection = LiquibaseUtils.getConnection(databaseSession.getDatabase());
            statement = connection.createStatement();
            statement.setFetchSize(fetchSize);
            resultSet = statement.executeQuery("SELECT * FROM " + escapedTableName + orderByClause(table, orderingColumn));
            while (resultSet.next() && !cancelState.isCanceled()) {
                forCurrentRow(table, resultSet, effect);
                numberOfRows++;
                if (numberOfRows % 10000 == 0) {
                    log.trace("{}: {} rows processed", tableName, numberOfRows);
                }
            }
            return numberOfRows;
        } catch (SQLException e) {
            throw new DataRetrievalFailureException("Could not fetch rows from " + tableName, e);
        } finally {
            JdbcUtils.closeResultSet(resultSet);
            JdbcUtils.closeStatement(statement);
        }
    }

    @Override
    @Nonnull
    public Iterable<String> getColumnNames(@Nonnull String tableName) {
        return chain(findTable(tableName).getColumns())
                .transform(TO_LOWERCASE_COLUMN_NAME)
                .toSet();
    }

    @Override
    @Nonnull
    public String getDatabaseType() {
        return databaseSession.getDatabase().getTypeName();
    }

    @Override
    @Nonnull
    public Iterable<String> getTableNames() {
        return chain(databaseSession.getSnapshot().getTables())
                .transform(TO_LOWERCASE_TABLE_NAME)
                .filter(IS_HIBERNATE_TEMPORARY_TABLE.negate())
                .toList();
    }

    /**
     * This implementation commits the change to the underlying database if the number of uncommitted changes has
     * reached the limit specified in the constructor.
     */
    @Override
    public void insert(@Nonnull InsertDataChange change) {
        List<ColumnConfig> columns = change.getColumns();

        //applyChange(change);
        Database database = databaseSession.getDatabase();
        StringBuilder builder = new StringBuilder("insert into ")
                .append(database.escapeTableName(change.getSchemaName(), change.getTableName()))
                .append(" (");

        StringBuilder questions = new StringBuilder();
        for (ColumnConfig column : columns) {
            if (questions.length() > 0) {
                builder.append(", ");
                questions.append(", ");
            }
            builder.append(database.escapeColumnName(change.getSchemaName(), change.getTableName(), column.getName()));
            questions.append("?");
        }
        builder.append(") values (").append(questions).append(")");

        Connection connection = LiquibaseUtils.getConnection(database);

        PreparedStatement statement = null;
        try {
            statement = connection.prepareStatement(builder.toString());

            for (int i = 0; i < columns.size(); ++i) {
                ColumnConfig column = columns.get(i);
                Object value = column.getValueObject();

                // Look for any instances of {@link DecodeBase64Function} which we use to mark any binary data that has
                // to be BASE64 encoded in Liquibase, but which the database driver can just process directly as a raw
                // byte[] array.
                if (value instanceof DecodeBase64Function) {
                    value = ((DecodeBase64Function) value).decode();
                } else if (value instanceof DatabaseFunction) {
                    // Got an unrecognized DatabaseFunction.  This should never happen as Liquibase should have
                    // evaluated them all already, and the database driver would have no idea what to do with one
                    // anyway.  But the "impossible" may sometimes happen, see
                    // https://support.atlassian.com/browse/SSP-9599.  In case this ever happens again, log a
                    // warning but don't interfere.
                    log.warn("Unrecognized DatabaseFunction {} for column \"{}\"", value, column.getName());
                }

                statement.setObject(i + 1, value);
            }
            statement.executeUpdate();
        } catch (SQLException e) {
            log.error("Unable to insert into table {} data {} with statement {}", change.getTableName(),
                    change.getColumns().stream()
                            .map(columnConfig -> columnConfig.getName() + "=" + columnConfig.getValueObject())
                            .collect(Collectors.toList()),
                    builder);
            throw new LiquibaseChangeExecutionException(change, e);
        } finally {
            JdbcUtils.closeStatement(statement);
        }

        commitIfBlockFilled();
    }

    /**
     * Sets the {@code ApplicationContext}, which is used to retrieve the "liquibasePrototype" when
     * {@link #createSchema(DataSource) creating the schema}.
     * <p>
     * Note: This is done using {@code ApplicationContextAware} and a setter to allow it to be used from the
     * backup/restore client, which does not use Spring annotations. Changing this will break restore client
     * processing.
     *
     * @param applicationContext the application context
     */
    @Override
    public void setApplicationContext(ApplicationContext applicationContext) {
        this.applicationContext = applicationContext;
    }

    @Override
    public void withLock(Effect<LiquibaseDao> effect) {
        lock();
        try {
            effect.apply(this);
        } finally {
            unlock();
        }
    }

    @VisibleForTesting
    void applyChange(Change change) {
        try {
            change.init();
            databaseSession.getDatabase().executeStatements(change, EMPTY_CHANGE_LOG, NO_VISITORS);
            log.debug(change.getConfirmationMessage());
            databaseSession.incrementChangeCount();
        } catch (LiquibaseException e) {
            throw new LiquibaseChangeExecutionException(change, e);
        }
    }

    /**
     * Commits the current transaction in the underlying database, and resets the change counter.
     */
    @VisibleForTesting
    void commit() {
        try {
            databaseSession.getDatabase().commit();
            databaseSession.resetChangeCount();
        } catch (DatabaseException e) {
            throw new LiquibaseDataAccessException("Failed to commit changes to database", databaseSession.getDatabase(), e);
        }
    }

    @VisibleForTesting
    LockService getLockService(Database database) {
        return LockService.getInstance(database);
    }

    /**
     * Rolls back the current transaction and resets the change counter.
     */
    @VisibleForTesting
    void rollback() {
        try {
            databaseSession.getDatabase().rollback();
            databaseSession.resetChangeCount();
        } catch (DatabaseException e) {
            throw new LiquibaseDataAccessException("Failed to rollback changes to database", databaseSession.getDatabase(), e);
        }
    }

    /**
     * Transforms column values for any database-specific mapping idiosyncrasies we need to account for.
     */
    @VisibleForTesting
    Object transformValue(Database database, Object value, int dataType, int columnSize) {
        // map an Oracle NUMBER(1) back to a boolean
        if (database instanceof OracleDatabase && dataType == Types.DECIMAL && columnSize == 1
                && (BigDecimal.ZERO.equals(value) || BigDecimal.ONE.equals(value))) {
            value = BigDecimal.ONE.equals(value);
        }

        return value;
    }

    /**
     * A predicate that tests whether the table passed to it has a given name, ignoring upper/lower case differences.
     *
     * @param tableName the name to compare actual table names against
     * @return a new predicate
     */
    private static Predicate<Table> byTableName(final String tableName) {
        return table -> table.getName().equalsIgnoreCase(tableName);
    }

    /**
     * Commit only if the number of changes is equal to the commit block size
     */
    private void commitIfBlockFilled() {
        if (commitBlockSize > 0 && databaseSession.getChangeCount() >= commitBlockSize) {
            commit();
        }
    }

    /**
     * Finds the given Liquibase table in the database snapshot.
     *
     * @param tableName the name of the table to find
     * @return a Liquibase table
     * @throws NoSuchElementException if the table is not found
     */
    private Table findTable(String tableName) {
        try {
            return find(databaseSession.getSnapshot().getTables(), byTableName(tableName)::test);
        } catch (NoSuchElementException e) {
            throw new NoSuchElementException(String.format("No table '%s' in database snapshot", tableName));
        }
    }

    /**
     * Applies the given effect to the current row in the given result set, with reference to the structure of
     * the given table definition
     *
     * @param table     the definition of the table from which the row was fetched
     * @param resultSet the result set containing the row to apply the effect to
     * @param effect    the effect to apply
     */
    private void forCurrentRow(Table table, ResultSet resultSet, Effect<Map<String, Object>> effect) throws SQLException {
        ResultSetMetaData metaData = resultSet.getMetaData();
        int columnCount = metaData.getColumnCount();
        Map<String, Object> columnValues = Maps.newHashMapWithExpectedSize(columnCount);
        for (int i = 1; i <= columnCount; i++) {
            String columnName = getColumnName(metaData, i);
            columnValues.put(columnName, getValue(resultSet, i, table.getDatabase(), table.getColumn(columnName)));
        }
        effect.apply(columnValues);
    }

    private String getColumnName(ResultSetMetaData metaData, int columnIndex) throws SQLException {
        return metaData.getColumnName(columnIndex).toLowerCase();
    }

    private Object getValue(ResultSet resultSet, int index, Database database, Column column) throws SQLException {
        return transformValue(database, JdbcUtils.getResultSetValue(resultSet, index), column.getDataType(), column.getColumnSize());
    }

    /**
     * This implementation locks the Liquibase change log table, which is updated by the restore process.
     */
    private void lock() {
        log.debug("Locking Liquibase");
        LockService service = getLockService(databaseSession.getDatabase());
        try {
            service.waitForLock();
            lockService = some(service);
            log.debug("Liquibase locked");
        } catch (LockException e) {
            throw new CannotAcquireLockException("Could not lock Liquibase change log", e);
        }
    }

    private String orderByClause(Table table, String column) {
        return column == null ? "" :
                String.format(" ORDER BY %s ASC",
                        databaseSession.getDatabase().escapeColumnName(
                                table.getSchema(),
                                table.getName(),
                                //Gets the column name with its exact letter casing (since escapeColumnName() will
                                //quote the column name depending on the name and database, requiring an exact match).
                                //For more info, see the overrides of AbstractDatabase.escapeDatabaseObject(String)
                                //for each database.
                                table.getColumn(column).getName()));
    }

    private void unlock() {
        log.debug("Unlocking Liquibase");
        lockService.foreach(RELEASE_LOCK);
        lockService = none();
        log.debug("Liquibase unlocked");
    }
}
