package com.atlassian.stash.internal.liquibase;

import com.google.common.base.Function;
import com.google.common.collect.Collections2;
import com.google.common.collect.Iterables;
import com.google.common.collect.Maps;
import liquibase.change.custom.CustomTaskChange;
import liquibase.change.custom.CustomTaskRollback;
import liquibase.database.Database;
import liquibase.exception.CustomChangeException;
import liquibase.exception.RollbackImpossibleException;
import liquibase.exception.UnsupportedChangeException;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.dao.DataAccessException;
import org.springframework.jdbc.core.JdbcTemplate;
import org.springframework.jdbc.core.RowCallbackHandler;

import java.sql.ResultSet;
import java.sql.SQLException;
import java.sql.Types;
import java.util.Arrays;
import java.util.List;
import java.util.Map;

/**
 * Ensures all sta_comment.root_id are set to the correct non-null values for all replies.
 *
 * @since 2.0
 */
public class FillInRootCommentsChange extends AbstractCustomChange implements CustomTaskChange, CustomTaskRollback {

    public static final int MAX_REPLY_NESTING = 50;

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

    @Override
    public void execute(Database database) throws CustomChangeException {
        //Anchor root replies with their roots first - we fan out from here
        Map<Long, Long> rootRepliesAndTheirRoots = getRootRepliesAndTheirRoots(database);
        log.debug("Setting roots on {} top-level replies", rootRepliesAndTheirRoots.size());
        setRoots(database, rootRepliesAndTheirRoots);

        int level = 0;
        Map<Long, Long> unrootedReplies = getUnrootedReplies(database);
        while (level < MAX_REPLY_NESTING && !unrootedReplies.isEmpty()) {
            log.debug("Setting roots on {} replies at level {}", unrootedReplies.size(), level + 1);
            setRoots(database, unrootedReplies);
            unrootedReplies = getUnrootedReplies(database);
            level++;
        }

        //It's unlikely there are 50 levels of replies, more likely is that we have somehow failed
        //to reduce the work left in each successive level so we need to bail now
        if (!unrootedReplies.isEmpty()) {
            throw new CustomChangeException(
                    String.format("Failed to set comment roots - too much comment nesting detected. Gave up at level %d", level + 1));
        }
    }

    @Override
    public String getConfirmationMessage() {
        return "Comment roots successfully set for all replies";
    }

    @Override
    public void rollback(Database database) throws CustomChangeException, UnsupportedChangeException, RollbackImpossibleException {
        try {
            JdbcTemplate template = LiquibaseUtils.getJdbcTemplate(database);
            template.update("UPDATE sta_comment SET root_id = NULL");
        } catch (DataAccessException e) {
            log.error("Failed to rollback rooting of replies", e);
            throw new CustomChangeException("Failed to rollback data in sta_comment", e);
        }
    }

    private Map<Long, Long> getRootRepliesAndTheirRoots(Database database) throws CustomChangeException {
        try {
            final Map<Long, Long> result = Maps.newHashMap();
            JdbcTemplate template = LiquibaseUtils.getJdbcTemplate(database);
            template.query("select c1.id as reply_id, c2.id as root_id from sta_comment c1 " +
                    "inner join sta_comment c2 on c1.parent_id = c2.id where c2.parent_id is NULL",
                    new RowCallbackHandler() {
                        @Override
                        public void processRow(ResultSet rs) throws SQLException {
                            long replyId = rs.getLong("reply_id");
                            long rootId = rs.getLong("root_id");
                            result.put(replyId, rootId);
                        }
                    });
            return result;
        } catch (DataAccessException e) {
            log.error("Failed to query root replies", e);
            throw new CustomChangeException("Failed to select data from sta_comment", e);
        }
    }

    private Map<Long, Long> getUnrootedReplies(Database database) throws CustomChangeException {
        try {
            final Map<Long, Long> result = Maps.newHashMap();
            JdbcTemplate template = LiquibaseUtils.getJdbcTemplate(database);
            template.query("select c1.id as reply_id, c2.root_id as root_id from sta_comment c1 " +
                    "inner join sta_comment c2 on c1.parent_id = c2.id where c1.root_id is NULL and c2.root_id is not NULL",
                    new RowCallbackHandler() {
                        @Override
                        public void processRow(ResultSet rs) throws SQLException {
                            long replyId = rs.getLong("reply_id");
                            long rootId = rs.getLong("root_id");
                            result.put(replyId, rootId);
                        }
                    });
            return result;
        } catch (DataAccessException e) {
            log.error("Failed to query unrooted replies", e);
            throw new CustomChangeException("Failed to select data from sta_comment", e);
        }
    }

    private void setRoots(Database database, Map<Long, Long> repliesAndTheirRoots) throws CustomChangeException {
        if (!repliesAndTheirRoots.isEmpty()) {
            try {
                JdbcTemplate template = LiquibaseUtils.getJdbcTemplate(database);
                Iterable<List<Object[]>> batches = Iterables.partition(Collections2.transform(repliesAndTheirRoots.entrySet(),
                        new Function<Map.Entry<Long, Long>, Object[]>() {
                            @Override
                            public Object[] apply(Map.Entry<Long, Long> entry) {
                                return new Object[]{entry.getValue(), entry.getKey()}; //Note the order swap
                            }
                        }), MAX_BATCH_SIZE);

                boolean debug = log.isDebugEnabled();
                int batchCount = 0;
                for (List<Object[]> batch : batches) {
                    if (debug) {
                        log.debug("Executing batch {} of setting roots for this level of reply", batchCount);
                    }
                    int[] results = template.batchUpdate("UPDATE sta_comment SET root_id = ? WHERE id = ?",
                            batch, new int[]{Types.BIGINT, Types.BIGINT});
                    if (debug) {
                        log.debug("Results of executing batch {} of setting roots for this level of reply: {}",
                                batchCount, Arrays.toString(results));
                    }
                    ++batchCount;
                }
            } catch (DataAccessException e) {
                log.error("Failed to set roots for replies", e);
                throw new CustomChangeException("Failed to update data in sta_comment", e);
            }
        }
    }
}
