package org.drools.core.command.runtime;

import org.drools.core.command.impl.GenericCommand;
import org.drools.core.command.impl.KnowledgeCommandContext;
import org.drools.persistence.TransactionManager;
import org.drools.persistence.TransactionSynchronization;
import org.kie.api.runtime.EnvironmentName;
import org.kie.api.runtime.KieSession;
import org.kie.internal.command.Context;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * KieSession cannot be disposed during Transaction!
 *
 * Solution:
 * Check if Transaction is running if so register a Synchronisation
 * that disposes session after transaction completion
 *
 * http://mswiderski.blogspot.de/2012/10/dispose-session-in-cmt-environment.html
 *
 */
public class CustomDisposeCommand implements GenericCommand<Void> {

    private static final long serialVersionUID = 1L;
    private final Logger logger = LoggerFactory.getLogger(getClass());

    @Override
    public Void execute(Context context) {
        final KieSession ksession = ((KnowledgeCommandContext) context).getKieSession();

        logger.debug("Trying to dispose KieSession (" + ksession + "). Checking for active transactions.");

        TransactionManager tm = (TransactionManager) ksession.getEnvironment().get(EnvironmentName.TRANSACTION_MANAGER);

        final int txStatus = tm.getStatus();

        logger.debug("TransactionStatus: " + getTxStatusName(txStatus));


        if (txStatus == TransactionManager.STATUS_NO_TRANSACTION) {

            logger.debug("No active transaction: disposing KieSession (" + ksession + ") directly");

            ksession.dispose();
            logger.debug("KieSession disposed");
        } else {
            try {

                logger.debug("Active transaction: registering KieSession (" + ksession + ") for dispose");

                tm.registerTransactionSynchronization(new TransactionSynchronization() {

                    @Override
                    public void beforeCompletion() {
                        // not used here

                    }

                    @Override
                    public void afterCompletion(int arg0) {
                        ksession.dispose();
                        logger.debug("KieSession disposed");

                    }

                });
            } catch (Exception e) {
                e.printStackTrace();
            }

        }

        return null;
    }

    private String getTxStatusName(int txStatus){
        switch (txStatus){
            case TransactionManager.STATUS_ACTIVE:
                return "STATUS_ACTIVE";
            case TransactionManager.STATUS_NO_TRANSACTION:
                return "STATUS_NO_TRANSACTION";
            case TransactionManager.STATUS_COMMITTED:
                return "STATUS_COMMITTED";
            case TransactionManager.STATUS_ROLLEDBACK:
                return "STATUS_ROLLEDBACK";
            case TransactionManager.STATUS_UNKNOWN:
                return "STATUS_UNKNOWN";
            default: return "UNDEFINED";
        }
    }
}