package com.atlassian.stash.internal.util;

import org.springframework.transaction.PlatformTransactionManager;
import org.springframework.transaction.TransactionDefinition;
import org.springframework.transaction.TransactionStatus;
import org.springframework.transaction.interceptor.DefaultTransactionAttribute;

/**
 * Handles batching of transactions when processing a large number of elements.
 */
public class TransactionBatcher {

    private final PlatformTransactionManager transactionManager;

    private final int batchSize;
    private final String identifier;

    private TransactionStatus tx;
    private int count = 0;

    /**
     * @param transactionManager
     * @param batchSize          number of elements to batch before creating a new transaction
     * @param identifier         a string just to identify this transaction
     */
    public TransactionBatcher(PlatformTransactionManager transactionManager, int batchSize, String identifier) {
        this.transactionManager = transactionManager;
        this.batchSize = batchSize;
        this.identifier = identifier;
    }

    /**
     * To be called for each iteration of a loop, which may result in a new transaction being created.
     *
     * @return true if the batch size was reached.
     */
    public boolean tick() {
        boolean ticked = ++count % batchSize == 0;
        if (ticked) {
            commit();
            start();
        }
        return ticked;
    }

    /**
     * Starts a new transaction for batching.
     */
    public void start() {
        DefaultTransactionAttribute def = new DefaultTransactionAttribute();
        def.setName("TransactionBatch:" + identifier + ":" + count);
        def.setPropagationBehavior(TransactionDefinition.PROPAGATION_REQUIRES_NEW);

        tx = transactionManager.getTransaction(def);
    }

    /**
     * Commit the current transaction.
     */
    public void commit() {
        transactionManager.commit(tx);
        tx = null;
    }

    /**
     * Rollback the current transaction.
     *
     * @return true if the transaction needed to be rolled back.
     */
    public boolean rollback() {
        if (tx != null && !tx.isCompleted()) {
            transactionManager.rollback(tx);
            return true;
        }
        return false;
    }

    /**
     * @return the number of total elements processed.
     */
    public int getCount() {
        return count;
    }

}
