package org.jboss.cache.interceptors;

import org.apache.commons.logging.LogFactory;
import org.jboss.cache.DataContainer;
import org.jboss.cache.Fqn;
import org.jboss.cache.InvocationContext;
import org.jboss.cache.Modification;
import org.jboss.cache.NodeSPI;
import org.jboss.cache.commands.AbstractVisitor;
import org.jboss.cache.commands.ReversibleCommand;
import org.jboss.cache.commands.VisitableCommand;
import org.jboss.cache.commands.tx.CommitCommand;
import org.jboss.cache.commands.tx.OptimisticPrepareCommand;
import org.jboss.cache.commands.tx.PrepareCommand;
import org.jboss.cache.commands.tx.RollbackCommand;
import org.jboss.cache.commands.write.ClearDataCommand;
import org.jboss.cache.commands.write.MoveCommand;
import org.jboss.cache.commands.write.PutDataMapCommand;
import org.jboss.cache.commands.write.PutForExternalReadCommand;
import org.jboss.cache.commands.write.PutKeyValueCommand;
import org.jboss.cache.commands.write.RemoveKeyCommand;
import org.jboss.cache.commands.write.RemoveNodeCommand;
import org.jboss.cache.config.CacheLoaderConfig;
import org.jboss.cache.factories.annotations.Inject;
import org.jboss.cache.factories.annotations.Start;
import org.jboss.cache.interceptors.base.SkipCheckChainedInterceptor;
import org.jboss.cache.loader.CacheLoader;
import org.jboss.cache.loader.CacheLoaderManager;
import org.jboss.cache.transaction.GlobalTransaction;
import org.jboss.cache.transaction.TransactionEntry;

import javax.transaction.SystemException;
import javax.transaction.Transaction;
import javax.transaction.TransactionManager;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;

/**
 * Writes modifications back to the store on the way out: stores modifications back
 * through the CacheLoader, either after each method call (no TXs), or at TX commit.
 *
 * @author Bela Ban
 * @version $Id: CacheStoreInterceptor.java 6772 2008-09-22 15:53:19Z manik.surtani@jboss.com $
 */
public class CacheStoreInterceptor extends SkipCheckChainedInterceptor implements CacheStoreInterceptorMBean
{
   private CacheLoaderConfig loaderConfig = null;
   private TransactionManager txMgr = null;
   private HashMap txStores = new HashMap();
   private Map<GlobalTransaction, Set<Fqn>> preparingTxs = new ConcurrentHashMap<GlobalTransaction, Set<Fqn>>();
   private long cacheStores = 0;
   private CacheLoader loader;
   private DataContainer dataContainer;
   private CacheLoaderManager loaderManager;

   public CacheStoreInterceptor()
   {
      log = LogFactory.getLog(getClass());
      trace = log.isTraceEnabled();
   }

   @Inject
   protected void init(DataContainer dataContainer, CacheLoaderManager loaderManager, TransactionManager txManager, CacheLoaderConfig clConfig)
   {
      // never inject a CacheLoader at this stage - only a CacheLoaderManager, since the CacheLoaderManager only creates a CacheLoader instance when it @Starts.
      this.loaderManager = loaderManager;
      this.loaderConfig = clConfig;
      txMgr = txManager;
      this.dataContainer = dataContainer;
   }

   @Start
   protected void start()
   {
      // this should only happen after the CacheLoaderManager has started, since the CacheLoaderManager only creates the CacheLoader instance in it's @Start method.
      loader = loaderManager.getCacheLoader();
   }

   /**
    * if this is a shared cache loader and the call is of remote origin, pass up the chain
    */
   @Override
   public boolean skipInterception(InvocationContext ctx, VisitableCommand command)
   {
      if (!ctx.isOriginLocal() && loaderConfig.isShared())
      {
         if (trace)
         {
            log.trace("Passing up method call and bypassing this interceptor since the cache loader is shared and this call originated remotely.");
         }
         return true;
      }
      return false;
   }

   @Override
   protected Object handleCommitCommand(InvocationContext ctx, CommitCommand command) throws Throwable
   {
      if (inTransaction())
      {
         if (ctx.isTxHasMods())
         {
            // this is a commit call.
            GlobalTransaction gtx = command.getGlobalTransaction();
            if (trace) log.trace("Calling loader.commit() for gtx " + gtx);
            // sync call (a write) on the loader
            // ignore modified FQNs
            // List fqnsModified = getFqnsFromModificationList(txTable.get(globalTransaction).getCacheLoaderModifications());
            try
            {
               loader.commit(gtx);
            }
            catch (Throwable t)
            {
               preparingTxs.remove(gtx);
               throw t;
            }
            if (getStatisticsEnabled())
            {
               Integer puts = (Integer) txStores.get(gtx);
               if (puts != null)
               {
                  cacheStores = cacheStores + puts;
               }
               txStores.remove(gtx);
            }
            Object returnValue = invokeNextInterceptor(ctx, command);
            // persist additional internal state, if any, and then clean up internal resources
            Set<Fqn> affectedFqns = preparingTxs.remove(gtx);
            if (affectedFqns != null)
            {
               storeInternalState(affectedFqns);
            }
            return returnValue;
         }
         else
         {
            if (trace) log.trace("Commit called with no modifications; ignoring.");
         }
      }
      return invokeNextInterceptor(ctx, command);
   }

   @Override
   protected Object handleRollbackCommand(InvocationContext ctx, RollbackCommand command) throws Throwable
   {
      if (inTransaction())
      {
         if (trace) log.trace("transactional so don't put stuff in the cloader yet.");
         if (ctx.isTxHasMods())
         {
            GlobalTransaction gtx = command.getGlobalTransaction();
            // this is a rollback method
            if (preparingTxs.containsKey(gtx))
            {
               preparingTxs.remove(gtx);
               loader.rollback(gtx);
            }
            if (getStatisticsEnabled()) txStores.remove(gtx);
         }
         else
         {
            if (trace) log.trace("Rollback called with no modifications; ignoring.");
         }
      }
      return invokeNextInterceptor(ctx, command);
   }

   @Override
   protected Object handleOptimisticPrepareCommand(InvocationContext ctx, OptimisticPrepareCommand command) throws Throwable
   {
      return handlePrepareCommand(ctx, command);
   }

   @Override
   protected Object handlePrepareCommand(InvocationContext ctx, PrepareCommand command) throws Throwable
   {
      if (inTransaction())
      {
         if (trace) log.trace("transactional so don't put stuff in the cloader yet.");
         prepareCacheLoader(command.getGlobalTransaction(), ctx.getTransactionEntry(), command.isOnePhaseCommit());
      }
      return invokeNextInterceptor(ctx, command);
   }

   /**
    * remove() methods need to be applied to the CacheLoader before passing up the call: a listener might
    * access an element just removed, causing the CacheLoader to *load* the element before *removing* it.
    */
   @Override
   protected Object handleRemoveNodeCommand(InvocationContext ctx, RemoveNodeCommand command) throws Throwable
   {
      if (!inTransaction())
      {
         loader.remove(command.getFqn());
      }
      return invokeNextInterceptor(ctx, command);
   }

   @Override
   protected Object handleRemoveKeyCommand(InvocationContext ctx, RemoveKeyCommand command) throws Throwable
   {
      if (!inTransaction())
      {
         Object returnValue = loader.remove(command.getFqn(), command.getKey());
         invokeNextInterceptor(ctx, command);
         return returnValue;
      }
      return invokeNextInterceptor(ctx, command);
   }

   @Override
   protected Object handleRemoveDataCommand(InvocationContext ctx, ClearDataCommand command) throws Throwable
   {
      if (!inTransaction())
      {
         loader.removeData(command.getFqn());
         // we need to mark this node as data loaded
         NodeSPI n = dataContainer.peek(command.getFqn(), false, false);
         if (n != null)
         {
            n.setDataLoaded(true);
         }
      }
      return invokeNextInterceptor(ctx, command);
   }

   @Override
   protected Object handleMoveCommand(InvocationContext ctx, MoveCommand command) throws Throwable
   {
      Object returnValue = invokeNextInterceptor(ctx, command);
      if (inTransaction())
      {
         return returnValue;
      }
      Fqn newNodeFqn = Fqn.fromRelativeElements(command.getTo(), command.getFqn().getLastElement());
      recursiveMove(command.getFqn(), newNodeFqn);
      loader.remove(command.getFqn());
      return returnValue;
   }

   @Override
   protected Object handlePutDataMapCommand(InvocationContext ctx, PutDataMapCommand command) throws Throwable
   {
      Object returnValue = invokeNextInterceptor(ctx, command);
      if (inTransaction())
      {
         return returnValue;
      }

      loader.put(command.getFqn(), command.getData());
      if (getStatisticsEnabled()) cacheStores++;

      return returnValue;
   }

   @Override
   protected Object handlePutKeyValueCommand(InvocationContext ctx, PutKeyValueCommand command) throws Throwable
   {
      Object returnValue = invokeNextInterceptor(ctx, command);
      if (inTransaction())
      {
         return returnValue;
      }
      returnValue = loader.put(command.getFqn(), command.getKey(), command.getValue());
      if (getStatisticsEnabled()) cacheStores++;

      return returnValue;
   }

   @Override
   protected Object handlePutForExternalReadCommand(InvocationContext ctx, PutForExternalReadCommand command) throws Throwable
   {
      return handlePutKeyValueCommand(ctx, command);
   }

   private boolean inTransaction() throws SystemException
   {
      return txMgr != null && txMgr.getTransaction() != null;
   }

   private void storeInternalState(Set<Fqn> affectedFqns) throws Exception
   {
      if (configuration.isNodeLockingOptimistic())
      {
         // we need to suspend any txs here since they would be in the tx-committed state and if the loader attempts to
         // use JTA (E.g., a JDBC CL using connections from a tx aware datasource) it will fail since the tx is in an
         // illegal state to perform writes.  See JBCACHE-1408.
         Transaction tx = txMgr.suspend();
         try
         {
            for (Fqn f : affectedFqns)
            {
               // NOT going to store tombstones!!
               NodeSPI n = dataContainer.peek(f, false, false);
               if (n != null)
               {
                  Map internalState = n.getInternalState(true);
                  loader.put(f, internalState);
               }
            }
         }
         finally
         {
            txMgr.resume(tx);
         }
      }
   }

   private void recursiveMove(Fqn fqn, Fqn newFqn) throws Exception
   {
      loader.put(newFqn, loader.get(fqn));
      //recurse
      Set childrenNames = loader.getChildrenNames(fqn);
      if (childrenNames != null)
      {
         for (Object child : childrenNames)
         {
            recursiveMove(Fqn.fromRelativeElements(fqn, child), Fqn.fromRelativeElements(newFqn, child));
         }
      }
   }

   public long getCacheLoaderStores()
   {
      return cacheStores;
   }

   @Override
   public void resetStatistics()
   {
      cacheStores = 0;
   }

   @Override
   public Map<String, Object> dumpStatistics()
   {
      Map<String, Object> retval = new HashMap<String, Object>();
      retval.put("CacheLoaderStores", cacheStores);
      return retval;
   }

   private void prepareCacheLoader(GlobalTransaction gtx, TransactionEntry entry, boolean onePhase) throws Throwable
   {
      if (entry == null)
      {
         throw new Exception("entry for transaction " + gtx + " not found in transaction table");
      }
      List<ReversibleCommand> modifications = entry.getModifications();
      if (modifications.size() == 0)
      {
         if (trace) log.trace("Transaction has not logged any modifications!");
         return;
      }
      if (trace) log.trace("Cache loader modification list: " + modifications);
      StoreModificationsBuilder modsBuilder = new StoreModificationsBuilder(getStatisticsEnabled());
      for (ReversibleCommand cacheCommand : modifications)
      {
         cacheCommand.acceptVisitor(null, modsBuilder);
      }
      if (trace)
      {
         log.trace("Converted method calls to cache loader modifications.  List size: " + modsBuilder.modifications.size());
      }
      if (modsBuilder.modifications.size() > 0)
      {
         loader.prepare(gtx, modsBuilder.modifications, onePhase);

         preparingTxs.put(gtx, modsBuilder.affectedFqns);
         if (getStatisticsEnabled() && modsBuilder.putCount > 0)
         {
            txStores.put(gtx, modsBuilder.putCount);
         }
      }
   }

   public static class StoreModificationsBuilder extends AbstractVisitor
   {
      boolean generateStatistics;
      int putCount;
      Set<Fqn> affectedFqns = new HashSet<Fqn>();
      List<Modification> modifications = new ArrayList<Modification>();

      public StoreModificationsBuilder(boolean generateStatistics)
      {
         this.generateStatistics = generateStatistics;
      }

      @Override
      public Object visitPutDataMapCommand(InvocationContext ctx, PutDataMapCommand command) throws Throwable
      {
         if (generateStatistics) putCount++;
         modifications.add(new Modification(Modification.ModificationType.PUT_DATA, command.getFqn(), command.getData()));
         affectedFqns.add(command.getFqn());
         return null;
      }

      @Override
      public Object visitPutKeyValueCommand(InvocationContext ctx, PutKeyValueCommand command) throws Throwable
      {
         if (generateStatistics) putCount++;
         modifications.add(new Modification(Modification.ModificationType.PUT_KEY_VALUE, command.getFqn(),
               command.getKey(), command.getValue()));
         affectedFqns.add(command.getFqn());
         return null;
      }

      @Override
      public Object visitRemoveKeyCommand(InvocationContext ctx, RemoveKeyCommand command) throws Throwable
      {
         modifications.add(new Modification(Modification.ModificationType.REMOVE_KEY_VALUE, command.getFqn(), command.getKey()));
         affectedFqns.add(command.getFqn());
         return null;
      }

      @Override
      public Object visitClearDataCommand(InvocationContext ctx, ClearDataCommand command) throws Throwable
      {
         modifications.add(new Modification(Modification.ModificationType.REMOVE_DATA, command.getFqn()));
         affectedFqns.add(command.getFqn());
         return null;
      }

      @Override
      public Object visitRemoveNodeCommand(InvocationContext ctx, RemoveNodeCommand command) throws Throwable
      {
         modifications.add(new Modification(Modification.ModificationType.REMOVE_NODE, command.getFqn()));
         affectedFqns.add(command.getFqn());
         return null;
      }

      @Override
      public Object visitMoveCommand(InvocationContext ctx, MoveCommand command) throws Throwable
      {
         Fqn moveFrom = command.getFqn();
         affectedFqns.add(command.getFqn());
         affectedFqns.add(moveFrom);
         Modification mod = new Modification(Modification.ModificationType.MOVE, moveFrom, command.getTo());
         modifications.add(mod);
         return null;
      }
   }
}
