package org.jboss.cache.util.internals;

import org.jboss.cache.Cache;
import org.jboss.cache.Fqn;
import org.jboss.cache.InvocationContext;
import org.jboss.cache.RPCManager;
import org.jboss.cache.commands.ReplicableCommand;
import org.jboss.cache.commands.remote.ReplicateCommand;
import org.jboss.cache.commands.tx.PrepareCommand;
import org.jboss.cache.factories.ComponentRegistry;
import org.jboss.cache.marshall.CommandAwareRpcDispatcher;
import org.jboss.cache.marshall.InactiveRegionAwareRpcDispatcher;
import org.jboss.cache.marshall.Marshaller;
import org.jboss.cache.marshall.RegionalizedMethodCall;
import org.jboss.cache.util.TestingUtil;
import org.jgroups.blocks.RpcDispatcher;
import org.jgroups.blocks.RpcDispatcher.Marshaller2;
import org.jgroups.util.Buffer;

import java.io.InputStream;
import java.io.ObjectInputStream;
import java.io.ObjectOutputStream;
import java.util.Arrays;
import java.util.HashSet;
import java.util.Iterator;
import java.util.Set;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;

/**
 * Utility class that notifies when certain commands were asynchronously replicated on secondary cache.
 * Especially useful for avaoiding Thread.sleep() statements.
 * <p/>
 * Usage:
 * <pre>
 *   Cache c1, c2; //these being two async caches
 *   AsyncReplicationListener listener2 = new AsyncReplicationListener(c2);
 *   listener2.expect(PutKeyValueCommand.class);
 *   c1.put(fqn, key, value);
 *   listener2.waitForReplicationToOccur(1000); // -this will block here untill c2 recieves the PutKeyValueCommand command
 * </pre>
 * Lifecycle - after being used (i.e. waitForReplicationToOccur returns sucessfully) the object returns to the
 * non-initialized state and *can* be reused through expect-wait cycle.
 * <b>Note</b>:  this class might be used aswell for sync caches, e.g. a test could have subclasses which use sync and
 * async replication
 *
 * @author Mircea.Markus@jboss.com
 * @since 2.2
 */
public class ReplicationListener
{
   private CountDownLatch latch = new CountDownLatch(1);
   private Set<Class<? extends ReplicableCommand>> expectedCommands;

   /**
    * Builds a listener that will observe the given cache for recieving replication commands.
    */
   public ReplicationListener(Cache cache)
   {
      ComponentRegistry componentRegistry = TestingUtil.extractComponentRegistry(cache);
      RPCManager rpcManager = componentRegistry.getComponent(RPCManager.class);
      CommandAwareRpcDispatcher realDispatcher = (CommandAwareRpcDispatcher) TestingUtil.extractField(rpcManager, "rpcDispatcher");
      RpcDispatcher.Marshaller2 realMarshaller = (Marshaller2) realDispatcher.getMarshaller();

      RpcDispatcher.Marshaller delegate = null;

      if (realDispatcher instanceof InactiveRegionAwareRpcDispatcher)
         delegate = new RegionMarshallerDelegate((Marshaller) realMarshaller);
      else
         delegate = new MarshallerDelegate(realMarshaller);

      realDispatcher.setMarshaller(delegate);
      realDispatcher.setRequestMarshaller(delegate);
      realDispatcher.setResponseMarshaller(delegate);
   }

   private class MarshallerDelegate implements RpcDispatcher.Marshaller
   {
      RpcDispatcher.Marshaller marshaller;

      private MarshallerDelegate(RpcDispatcher.Marshaller marshaller)
      {
         this.marshaller = marshaller;
      }

      public byte[] objectToByteBuffer(Object obj) throws Exception
      {
         return marshaller.objectToByteBuffer(obj);
      }

      public Object objectFromByteBuffer(byte bytes[]) throws Exception
      {
         Object result = marshaller.objectFromByteBuffer(bytes);
         if (result instanceof ReplicateCommand && expectedCommands != null)
         {
            ReplicateCommand replicateCommand = (ReplicateCommand) result;
            return new ReplicateCommandDelegate(replicateCommand);
         }
         return result;
      }
   }

   /**
    * We want the notification to be performed only *after* the remote command is executed.
    */
   private class ReplicateCommandDelegate extends ReplicateCommand
   {
      ReplicateCommand realOne;

      private ReplicateCommandDelegate(ReplicateCommand realOne)
      {
         this.realOne = realOne;
      }

      @Override
      public Object perform(InvocationContext ctx) throws Throwable
      {
         try
         {
            return realOne.perform(ctx);
         }
         finally
         {
            System.out.println("Processed command: " + realOne);
            Iterator<Class<? extends ReplicableCommand>> it = expectedCommands.iterator();
            while (it.hasNext())
            {
               Class<? extends ReplicableCommand> replicableCommandClass = it.next();
               if (realOne.containsCommandType(replicableCommandClass))
               {
                  it.remove();
               }
               else if (realOne.getSingleModification() instanceof PrepareCommand) //explicit transaction
               {
                  PrepareCommand prepareCommand = (PrepareCommand) realOne.getSingleModification();
                  if (prepareCommand.containsModificationType(replicableCommandClass))
                  {
                     it.remove();
                  }
               }
            }
            if (expectedCommands.isEmpty())
            {
               latch.countDown();
            }
         }
      }
   }

   /**
    * Needed for region based marshalling.
    */
   private class RegionMarshallerDelegate implements Marshaller
   {
      private Marshaller realOne;

      private RegionMarshallerDelegate(Marshaller realOne)
      {
         this.realOne = realOne;
      }

      public void objectToObjectStream(Object obj, ObjectOutputStream out) throws Exception
      {
         realOne.objectToObjectStream(obj, out);
      }

      public Object objectFromObjectStream(ObjectInputStream in) throws Exception
      {
         return realOne.objectFromObjectStream(in);
      }

      public Object objectFromStream(InputStream is) throws Exception
      {
         return realOne.objectFromStream(is);
      }

      public void objectToObjectStream(Object obj, ObjectOutputStream out, Fqn region) throws Exception
      {
         realOne.objectToObjectStream(obj, out, region);
      }

      public RegionalizedMethodCall regionalizedMethodCallFromByteBuffer(byte[] buffer) throws Exception
      {
         RegionalizedMethodCall result = realOne.regionalizedMethodCallFromByteBuffer(buffer);
         if (result.command instanceof ReplicateCommand && expectedCommands != null)
         {
            ReplicateCommand replicateCommand = (ReplicateCommand) result.command;
            result.command = new ReplicateCommandDelegate(replicateCommand);
         }
         return result;
      }

      public RegionalizedMethodCall regionalizedMethodCallFromObjectStream(ObjectInputStream in) throws Exception
      {
         return realOne.regionalizedMethodCallFromObjectStream(in);
      }

      public byte[] objectToByteBuffer(Object o) throws Exception
      {
         return realOne.objectToByteBuffer(o);
      }

      public Object objectFromByteBuffer(byte[] bytes) throws Exception
      {
         return realOne.objectFromByteBuffer(bytes);
      }

      public Buffer objectToBuffer(Object obj) throws Exception
      {
         return realOne.objectToBuffer(obj);
      }

      public Object objectFromByteBuffer(byte[] buf, int offset, int length) throws Exception
      {
         return realOne.objectFromByteBuffer(buf, offset, length);
      }
   }

   /**
    * Blocks for the elements specified through {@link #expect(Class[])} invocations to be replicated in this cache.
    * if replication does not occur in the give timeout then an exception is being thrown.
    */
   public void waitForReplicationToOccur(long timeoutMillis)
   {
      System.out.println("enter... ReplicationListener.waitForReplicationToOccur");
      waitForReplicationToOccur(timeoutMillis, TimeUnit.MILLISECONDS);
      System.out.println("exit... ReplicationListener.waitForReplicationToOccur");
   }

   /**
    * Similar to {@link #waitForReplicationToOccur(long)} except that this method provides more flexibility in time units.
    *
    * @param timeout  the maximum time to wait
    * @param timeUnit the time unit of the <tt>timeout</tt> argument.
    */
   public void waitForReplicationToOccur(long timeout, TimeUnit timeUnit)
   {
      assert expectedCommands != null : "there are no replication expectations; please use AsyncReplicationListener.expect(...) before calling this method";
      try
      {
         if (!latch.await(timeout, timeUnit))
         {
            assert false : "waiting for more than " + timeout + " " + timeUnit + " and following commands did not replicate: " + expectedCommands;
         }
      }
      catch (InterruptedException e)
      {
         throw new IllegalStateException("unexpected", e);
      }
      finally
      {
         expectedCommands = null;
         latch = new CountDownLatch(1);
      }
   }

   /**
    * {@link #waitForReplicationToOccur(long)} will block untill all the commands specified here are being replicated
    * to this cache. The method can be called several times with various arguments.
    */
   public void expect(Class<? extends ReplicableCommand>... expectedCommands)
   {
      if (this.expectedCommands == null)
      {
         this.expectedCommands = new HashSet<Class<? extends ReplicableCommand>>();
      }
      this.expectedCommands.addAll(Arrays.asList(expectedCommands));
   }

   /**
    * Waits untill first command is replicated.
    */
   public void expectAny()
   {
      expect();
   }
}
