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

import com.atlassian.bitbucket.util.CancelState;
import com.atlassian.fugue.Effect;
import com.atlassian.security.xml.SecureXmlParserFactory;
import com.atlassian.stash.internal.backup.DatabaseTable;
import com.atlassian.stash.internal.backup.liquibase.processors.ChangeSetProcessor;
import com.atlassian.stash.internal.backup.liquibase.processors.CompositeChangeSetProcessor;
import com.atlassian.stash.internal.backup.liquibase.processors.NullFilteringChangeSetProcessor;
import com.atlassian.stash.internal.backup.liquibase.xml.DefaultXmlEncoder;
import com.google.common.collect.Iterables;
import com.google.common.io.ByteStreams;
import com.google.common.io.Closeables;
import com.google.common.io.Flushables;
import org.apache.commons.lang3.StringUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;
import org.xml.sax.*;

import javax.xml.stream.XMLStreamException;
import java.io.*;

import static com.atlassian.security.xml.SecureXmlParserFactory.emptyEntityResolver;
import static com.atlassian.stash.internal.backup.DatabaseTableAttribute.PREPOPULATED;
import static com.atlassian.stash.internal.backup.liquibase.LiquibaseConstants.ENCODING;

/**
 * Default implementation for {@link LiquibaseMigrationDao}.
 */
@Component
public class DefaultLiquibaseMigrationDao implements LiquibaseMigrationDao {

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

    private final LiquibaseDao dao;
    /** factory generating the {@link LiquibaseXmlWriter} to write the Liquibase changelogs. */
    private final LiquibaseXmlWriterFactory writerFactory;

    @Autowired
    public DefaultLiquibaseMigrationDao(LiquibaseDao dao, LiquibaseXmlWriterFactory writerFactory) {
        this.writerFactory = writerFactory;
        this.dao = dao;
    }

    // --------- backup ---------

    /*
     * Implementation note: to reduce the memory footprint, the output is 'streamed' to the output stream; that is,
     * writes occur as rows are read from the underlying database, rather than being read into an object graph first.
     */
    @Override
    public void backup(OutputStream stream, String author, LiquibaseBackupMonitor monitor,
                       CancelState cancelState)  throws LiquibaseDataAccessException {
        log.info("Backup process started");

        LiquibaseXmlWriter writer = writerFactory.create(stream, author);
        try {
            // notify we started
            monitor.started(countTotalRows());

            // write headers
            writer.writeStartDocument(ENCODING, "1.0");
            writer.writeDatabaseChangeLogStartElement();

            // serialize each table as a Liquibase changeset
            for (String prePopulatedTableName : DatabaseTable.getTableNames(PREPOPULATED)) {
                if (cancelState.isCanceled()) {
                    return;
                }
                writer.writeChangeSetToDeleteRowsFromTable(prePopulatedTableName);
            }
            writeXmlForTables(monitor, writer, DatabaseTable.getTables(), cancelState);

            // finalise the backup
            writer.writeEndDocument();
            writer.flush();
            log.info("Backup process completed");

        } catch (XMLStreamException e) {
            throw new LiquibaseDataAccessException("An error occurred while writing to the output stream", e);
        } finally {
            // release the resources associated with the writer (which does not close the underlying output stream)
            try {
                writer.close();
            } catch (XMLStreamException e) {
                log.error("An error occurred while closing the output stream", e);
            }
            // close the database session
            dao.close();
        }
    }

    private long countTotalRows() {
        long totalRows = 0;
        for (DatabaseTable table : DatabaseTable.getTables()) {
            totalRows += dao.countRows(table.getTableName());
        }
        return totalRows;
    }

    /**
     * Writes a Liquibase changelog (conforming to the Liquibase change log DTD) for each table.
     */
    private void writeXmlForTables(LiquibaseBackupMonitor backupMonitor, LiquibaseXmlWriter writer,
                                   Iterable<DatabaseTable> tables, CancelState cancelState) throws XMLStreamException {
        int tableCount = Iterables.size(tables);
        long rowsWritten = 0;

        log.info("There are {} tables to back up", tableCount);
        for (DatabaseTable table : tables) {
            if (cancelState.isCanceled()) {
                return;
            }
            String tableName = table.getTableName();
            String orderingColumn = table.getOrderingColumn();
            log.info("Backing up {} table", tableName);

            // write all the table rows
            long numberOfRows = writeXmlForTable(dao, writer, tableName, orderingColumn, backupMonitor, cancelState);
            rowsWritten += numberOfRows;
            log.info("{} rows from the {} table were backed up", numberOfRows, tableName);
        }

        log.info("A total of {} rows in {} tables were backed up", rowsWritten, tableCount);
    }

    /**
     * Fetches the rows of the given table and serialises them as XML (conforming to the Liquibase change log DTD).
     */
    private long writeXmlForTable(LiquibaseDao dao, LiquibaseXmlWriter writer, String tableName,
                                  String orderingColumn, LiquibaseBackupMonitor monitor,
                                  CancelState cancelState) throws XMLStreamException {
        writer.writeChangeSetStartElement(tableName);
        Iterable<String> columnNames = dao.getColumnNames(tableName);
        long numberOfRows = dao.forEachRow(tableName, orderingColumn, cancelState,
                new WriteXmlForRowEffect(writer, tableName, columnNames, monitor));
        writer.writeEndElement(); // end of change set
        return numberOfRows;
    }

    // --------- restore ---------

    /*
     * Implementation notes:
     * - this reads the contents of the given input stream in two phases:
     *    - the first is a quick scan to obtain the changeset summaries (which include computing their relative 'weights'),
     *    - the second causes changes to be executed on a database connection;
     * - all referential integrity constraint checking is turned off while changes are being made. This method does
     *   not guarantee that constraints are retrospectively checked when constraint-checking is turned on again.
     */
    @Override
    public void restore(InputStream stream, File tempDir, final LiquibaseRestoreMonitor monitor,
                        final CancelState cancelState) {
        log.info("Restore process started");

        try {
            // create a temporary file to hold the contents of the input stream and then restore the database from it
            final File tempFile = createTempFile(tempDir);
            try {
                log.info("Examining backup data");
                log.debug("Saving stream contents to temporary file : {}", tempFile.getAbsolutePath());
                saveInputToFile(stream, tempFile);

                log.debug("Performing first pass through temporary file");
                final ChangeLogOutline changeSetMetaData = scan(tempFile);
                log.info("Found {} non-empty change sets to apply to the target database",
                        changeSetMetaData.nonEmptyChangeSetCount());

                log.info("Connecting to target database");
                dao.withLock(new Effect<LiquibaseDao>() {
                    @Override
                    public void apply(LiquibaseDao dao) {
                        log.debug("Beginning second pass through temporary file");
                        ChangeSetProcessor processor = createProcessor(monitor, changeSetMetaData);
                        ChangeSetReader handler = new ChangeSetReader(cancelState, processor, new DefaultXmlEncoder());
                        parse(tempFile, buildXmlReader(handler));
                        log.info("Restore process completed");
                    }
                });
            } finally {
                log.debug("Deleting temporary file");
                deleteFile(tempFile);
            }
        } finally {
            dao.close();
        }
    }

    /**
     * Create a processor for the Liquibase changesets read from the backup.
     */
    private ChangeSetProcessor createProcessor(LiquibaseRestoreMonitor monitor, ChangeLogOutline changeSetMetaData) {
        DatabaseUpdater databaseUpdater = new DatabaseUpdater(dao, monitor, changeSetMetaData);
        boolean isPostgres = StringUtils.containsIgnoreCase(dao.getDatabaseType(), "postgres");
        if (isPostgres) {
            return new CompositeChangeSetProcessor(new NullFilteringChangeSetProcessor(), databaseUpdater);
        }
        return databaseUpdater;
    }

    /**
     * Scans the given XML file to find out how many change sets there are, and how many changes are in each change set.
     * This is used later for progress reporting.
     *
     * @param xmlFile the file in which the change sets are specified
     * @return an iterable of change set meta data
     */
    private ChangeLogOutline scan(File xmlFile) {
        ChangeSetScanner scanner = new ChangeSetScanner();
        parse(xmlFile, buildXmlReader(scanner));
        return new ChangeLogOutline(scanner.getChangeCounts());
    }

    private void parse(File xmlFile, XMLReader xmlReader) {
        InputStream inputStream = null;
        try {
            inputStream = new FileInputStream(xmlFile);
            InputSource inputSource = new InputSource(inputStream);
            inputSource.setEncoding(ENCODING);
            xmlReader.parse(inputSource);
        } catch (FileNotFoundException e) {
            throw new LiquibaseDataAccessException("Failed to open backup file in order to parse its contents", e);
        } catch (IOException e) {
            throw new LiquibaseDataAccessException("An error occurred while parsing the backup file", e);
        } catch (SAXException e) {
            throw new LiquibaseDataAccessException("SAX parsing error while parsing backup file", e);
        } finally {
            Closeables.closeQuietly(inputStream);
        }
    }

    /**
     * Creates a temporary file with prefix &quot;stash_&quot; and a &quot;.xml&quot; extension.
     * @param directory directory where to create the file
     */
    private File createTempFile(File directory) {
        try {
            return File.createTempFile("stash_", ".xml", directory);
        } catch (IOException e) {
            throw new LiquibaseDataAccessException("Failed to create a temporary file to hold the change log", e);
        }
    }

    private void deleteFile(File file) {
        if (file != null) {
            if (!file.delete()) {
                file.deleteOnExit();
            }
        }
    }

    /**
     * Saves the contents of the given input stream into the given destination file, which is overwritten if it exists.
     * <p>
     * <i>Note: This method does not close the input stream, but does close the output stream that it uses to write to the
     * destination file.</i>
     *
     * @param inputStream a stream of data to be saved
     * @param destination a file into which the data will be written. The file may, but needs not, exist when this
     *                    method is called, and it will exist once this method returns.
     */
    private void saveInputToFile(InputStream inputStream, File destination) {
        try (FileOutputStream to = new FileOutputStream(destination)) {
            ByteStreams.copy(inputStream, to);
            Flushables.flushQuietly(to);
        } catch (IOException e) {
            throw new LiquibaseDataAccessException("Failed to save the contents of the input stream to a temporary file", e);
        }
    }

    private XMLReader buildXmlReader(ContentHandler handler) {
        try {
            XMLReader xmlReader = SecureXmlParserFactory.newNamespaceAwareXmlReader();
            xmlReader.setErrorHandler(errorHandler);
            xmlReader.setContentHandler(handler);
            xmlReader.setEntityResolver(emptyEntityResolver());
            xmlReader.setProperty("http://java.sun.com/xml/jaxp/properties/schemaLanguage",
                    "http://www.w3.org/2001/XMLSchema");
            return xmlReader;
        } catch (SAXNotRecognizedException e) {
            throw new LiquibaseDataAccessException("Could not associate the XML schema language with the backup file", e);
        } catch (SAXNotSupportedException e) {
            throw new LiquibaseDataAccessException("Could not configure the XML schema language for parsing the backup file", e);
        }
    }

    private final ErrorHandler errorHandler = new ErrorHandler() {
        public void warning(SAXParseException exception) throws SAXException {
            throw exception;
        }

        public void error(SAXParseException exception) throws SAXException {
            throw exception;
        }

        public void fatalError(SAXParseException exception) throws SAXException {
            throw exception;
        }
    };
}
