package org.infinispan.statetransfer;

import static org.infinispan.test.TestingUtil.wrapInboundInvocationHandler;
import static org.infinispan.test.TestingUtil.wrapComponent;
import static org.testng.AssertJUnit.assertEquals;

import java.util.Map;
import java.util.concurrent.Future;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;

import org.infinispan.commands.ReplicableCommand;
import org.infinispan.commands.remote.CacheRpcCommand;
import org.infinispan.commands.tx.PrepareCommand;
import org.infinispan.commons.CacheException;
import org.infinispan.configuration.cache.CacheMode;
import org.infinispan.configuration.cache.ConfigurationBuilder;
import org.infinispan.distribution.MagicKey;
import org.infinispan.remoting.inboundhandler.AbstractDelegatingHandler;
import org.infinispan.remoting.inboundhandler.DeliverOrder;
import org.infinispan.remoting.inboundhandler.PerCacheInboundInvocationHandler;
import org.infinispan.remoting.inboundhandler.Reply;
import org.infinispan.remoting.responses.Response;
import org.infinispan.remoting.responses.UnsureResponse;
import org.infinispan.remoting.rpc.RpcManager;
import org.infinispan.remoting.transport.Address;
import org.infinispan.test.MultipleCacheManagersTest;
import org.infinispan.test.concurrent.StateSequencer;
import org.infinispan.transaction.lookup.DummyTransactionManagerLookup;
import org.infinispan.util.AbstractControlledRpcManager;
import org.infinispan.util.ControlledConsistentHashFactory;
import org.infinispan.util.logging.Log;
import org.infinispan.util.logging.LogFactory;
import org.testng.annotations.Test;

/**
 * Test for https://issues.jboss.org/browse/ISPN-6047
 *
 * @author Pedro Ruivo
 * @since 8.2
 */
@Test(groups = "functional", testName = "statetransfer.TxReplay3Test")
public class TxReplay3Test extends MultipleCacheManagersTest {

   private static final String VALUE_1 = "v1";
   private static final String VALUE_2 = "v2";

   private static final String TX1_LOCKED = "tx1:acquired_lock";
   private static final String TX1_UNSURE = "tx1:unsure_response";
   private static final String TX2_PENDING = "tx2:waiting_tx1";
   private static final String MAIN_ADVANCE = "main:advance";
   private static final String JOIN_NEW_NODE = "join:add_new_node";

   public void testReplay() throws Exception {
      final Object key = new MagicKey("TxReplay3Test", cache(0));
      final StateSequencer sequencer = new StateSequencer();
      sequencer.logicalThread("tx1", TX1_LOCKED, TX1_UNSURE);
      sequencer.logicalThread("tx2", TX2_PENDING);
      sequencer.logicalThread("join", JOIN_NEW_NODE);
      sequencer.logicalThread("main", MAIN_ADVANCE);
      sequencer.order(TX1_LOCKED, MAIN_ADVANCE, TX2_PENDING, JOIN_NEW_NODE, TX1_UNSURE);


      wrapComponent(cache(1), RpcManager.class,
                    (wrapOn, current) -> new UnsureResponseRpcManager(current, sequencer), true);
      Handler handler = wrapInboundInvocationHandler(cache(0), current -> new Handler(current, sequencer));
      handler.setOrigin(address(cache(2)));

      Future<Void> tx1 = fork(() -> {
         cache(1).put(key, VALUE_1);
         return null;
      });

      sequencer.advance(MAIN_ADVANCE);

      Future<Void> tx2 = fork(() -> {
         cache(2).put(key, VALUE_2);
         return null;
      });

      sequencer.enter(JOIN_NEW_NODE);
      addClusterEnabledCacheManager(config()).getCache();
      waitForClusterToForm();
      sequencer.exit(JOIN_NEW_NODE);


      tx1.get(30, TimeUnit.SECONDS);
      tx2.get(30, TimeUnit.SECONDS);

      assertEquals(VALUE_2, cache(0).get(key));
   }

   @Override
   protected void createCacheManagers() throws Throwable {
      createClusteredCaches(3, config());
   }

   private static ConfigurationBuilder config() {
      ConfigurationBuilder builder = getDefaultClusteredCacheConfig(CacheMode.DIST_SYNC, true);
      builder.transaction()
            .useSynchronization(false)
            .transactionManagerLookup(new DummyTransactionManagerLookup())
            .recovery().disable();
      builder.locking().lockAcquisitionTimeout(1, TimeUnit.MINUTES);
      builder.clustering()
            .remoteTimeout(1, TimeUnit.MINUTES)
            .hash().numOwners(1).numSegments(1)
            .consistentHashFactory(new ControlledConsistentHashFactory(0))
            .stateTransfer().fetchInMemoryState(false);
      return builder;
   }

   private static class UnsureResponseRpcManager extends AbstractControlledRpcManager {

      private static final Log log = LogFactory.getLog(UnsureResponseRpcManager.class);
      private final StateSequencer sequencer;
      private volatile boolean triggered = false;

      public UnsureResponseRpcManager(RpcManager realOne, StateSequencer sequencer) {
         super(realOne);
         this.sequencer = sequencer;
      }

      @Override
      protected Object beforeInvokeRemotely(ReplicableCommand command) {
         Object arg = super.beforeInvokeRemotely(command);
         log.debugf("Before invoke remotely %s", command);
         return arg;
      }

      @Override
      protected Map<Address, Response> afterInvokeRemotely(ReplicableCommand command, Map<Address, Response> responseMap, Object argument) {
         Map<Address, Response> result = super.afterInvokeRemotely(command, responseMap, argument);
         log.debugf("After invoke remotely %s. Responses=%s", command, result);
         if (!triggered && command instanceof PrepareCommand) {
            log.debugf("Triggering %s and %s", TX1_LOCKED, TX1_UNSURE);
            triggered = true;
            try {
               sequencer.advance(TX1_LOCKED);
               sequencer.advance(TX1_UNSURE);
            } catch (TimeoutException | InterruptedException e) {
               throw new CacheException(e);
            }
            for (Map.Entry<Address, Response> entry : result.entrySet()) {
               entry.setValue(UnsureResponse.INSTANCE);
            }
            log.debugf("After invoke remotely %s. New Responses=%s", command, result);
         }
         return result;
      }
   }

   private static class Handler extends AbstractDelegatingHandler {

      private static final Log log = LogFactory.getLog(Handler.class);
      private final StateSequencer sequencer;
      private volatile boolean triggered = false;
      private volatile Address origin;

      public Handler(PerCacheInboundInvocationHandler delegate, StateSequencer sequencer) {
         super(delegate);
         this.sequencer = sequencer;
      }

      public void setOrigin(Address origin) {
         this.origin = origin;
      }

      @Override
      protected boolean beforeHandle(CacheRpcCommand command, Reply reply, DeliverOrder order) {
         log.debugf("Before invoking %s. expected origin=%s", command, origin);
         return super.beforeHandle(command, reply, order);
      }

      @Override
      protected void afterHandle(CacheRpcCommand command, DeliverOrder order, boolean delegated) {
         super.afterHandle(command, order, delegated);
         log.debugf("After invoking %s. expected origin=%s", command, origin);
         if (!triggered && command instanceof PrepareCommand && command.getOrigin().equals(origin)) {
            log.debugf("Triggering %s.", TX2_PENDING);
            triggered = true;
            try {
               sequencer.advance(TX2_PENDING);
            } catch (TimeoutException | InterruptedException e) {
               throw new CacheException(e);
            }
         }
      }
   }
}
