package org.infinispan.distribution;

import org.infinispan.Cache;
import org.infinispan.commands.VisitableCommand;
import org.infinispan.commands.write.PutKeyValueCommand;
import org.infinispan.configuration.cache.CacheMode;
import org.infinispan.configuration.cache.ConfigurationBuilder;
import org.infinispan.container.DataContainer;
import org.infinispan.container.entries.ImmortalCacheEntry;
import org.infinispan.container.entries.InternalCacheEntry;
import org.infinispan.container.entries.MortalCacheEntry;
import org.infinispan.distribution.group.Grouper;
import org.infinispan.distribution.ch.ConsistentHash;
import org.infinispan.interceptors.InterceptorChain;
import org.infinispan.interceptors.base.CommandInterceptor;
import org.infinispan.manager.EmbeddedCacheManager;
import org.infinispan.remoting.transport.Address;
import org.infinispan.test.MultipleCacheManagersTest;
import org.infinispan.test.TestingUtil;
import org.infinispan.transaction.LockingMode;
import org.infinispan.test.fwk.TransportFlags;
import org.infinispan.util.concurrent.IsolationLevel;

import javax.transaction.TransactionManager;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.concurrent.TimeUnit;


public abstract class BaseDistFunctionalTest<K, V> extends MultipleCacheManagersTest {
   protected String cacheName;
   protected int INIT_CLUSTER_SIZE = 4;
   protected Cache<K, V> c1 = null, c2 = null, c3 = null, c4 = null;
   protected ConfigurationBuilder configuration;
   protected List<Cache<K, V>> caches;
   protected List<Address> cacheAddresses;
   protected boolean sync = true;
   protected boolean tx = false;
   protected boolean testRetVals = true;
   protected boolean l1CacheEnabled = true;
   protected boolean l1OnRehash = false;
   protected int l1Threshold = 5;
   protected boolean performRehashing = false;
   protected boolean batchingEnabled = false;
   protected int numOwners = 2;
   protected int lockTimeout = 45;
   protected boolean groupsEnabled = false;
   protected List<Grouper<?>> groupers;
   protected LockingMode lockingMode;

   protected void createCacheManagers() throws Throwable {
      cacheName = "dist";
      configuration = buildConfiguration();
      // Create clustered caches with failure detection protocols on
      caches = createClusteredCaches(INIT_CLUSTER_SIZE, cacheName, configuration,
                                     new TransportFlags().withFD(false));

      if (INIT_CLUSTER_SIZE > 0) c1 = caches.get(0);
      if (INIT_CLUSTER_SIZE > 1) c2 = caches.get(1);
      if (INIT_CLUSTER_SIZE > 2) c3 = caches.get(2);
      if (INIT_CLUSTER_SIZE > 3) c4 = caches.get(3);

      cacheAddresses = new ArrayList<Address>(INIT_CLUSTER_SIZE);
      for (Cache cache : caches) {
         EmbeddedCacheManager cacheManager = cache.getCacheManager();
         cacheAddresses.add(cacheManager.getAddress());
      }
   }

   protected ConfigurationBuilder buildConfiguration() {
      ConfigurationBuilder configuration = getDefaultClusteredCacheConfig(sync ? CacheMode.DIST_SYNC : CacheMode.DIST_ASYNC, tx);
      configuration.clustering().stateTransfer().fetchInMemoryState(performRehashing);
      if (lockingMode != null) {
         configuration.transaction().lockingMode(lockingMode);
      }
      configuration.clustering().hash().numOwners(numOwners);
      if (!testRetVals) {
         configuration.unsafe().unreliableReturnValues(true);
         // we also need to use repeatable read for tests to work when we dont have reliable return values, since the
         // tests repeatedly queries changes
         configuration.locking().isolationLevel(IsolationLevel.REPEATABLE_READ);
      }
      if (tx) {
         configuration.invocationBatching().enable();
      }
      if (sync) configuration.clustering().sync().replTimeout(60, TimeUnit.SECONDS);
      configuration.locking().lockAcquisitionTimeout(lockTimeout, TimeUnit.SECONDS);
      configuration.clustering().l1().enabled(l1CacheEnabled);
      if (groupsEnabled) {
          configuration.clustering().hash().groups().enabled(true);
          configuration.clustering().hash().groups().withGroupers(groupers);
      }
      if (l1CacheEnabled) configuration.clustering().l1().onRehash(l1OnRehash).invalidationThreshold(l1Threshold);
      return configuration;
   }

   // ----------------- HELPERS ----------------

   protected void initAndTest() {
      for (Cache<K, V> c : caches) assert c.isEmpty();

      // TODO: A bit hacky, this should be moved somewhere else really...
      Cache<Object, Object> firstCache = (Cache<Object, Object>) caches.get(0);
      firstCache.put("k1", "value");
      asyncWait("k1", PutKeyValueCommand.class);
      assertOnAllCachesAndOwnership("k1", "value");
   }

   protected Address addressOf(Cache<?, ?> cache) {
      return DistributionTestHelper.addressOf(cache);
   }

   protected Cache<K, V> getFirstNonOwner(Object key) {
      return getNonOwners(key)[0];
   }
   
   protected Cache<K, V> getFirstOwner(Object key) {
      return getOwners(key)[0];
   }

   protected Cache<K, V> getSecondNonOwner(String key) {
      return getNonOwners(key)[1];
   }

   protected void assertOnAllCachesAndOwnership(Object key, String value) {
      assertOwnershipAndNonOwnership(key, l1CacheEnabled);
      // checking the values will bring the keys to L1, so we want to do it after checking ownership
      assertOnAllCaches(key, value);
   }

   protected void assertRemovedOnAllCaches(Object key) {
      assertOnAllCaches(key, null);
   }

   protected void assertOnAllCaches(Object key, String value) {
      for (Cache<K, V> c : caches) {
         Object realVal = c.get(key);
         if (value == null) {
            assert realVal == null : "Expecting [" + key + "] to equal [" + value + "] on cache ["
                  + addressOf(c) + "] but was [" + realVal + "]. Owners are " + Arrays.toString(getOwners(key));
         } else {
            assert value.equals(realVal) : "Expecting [" + key + "] to equal [" + value + "] on cache ["
                  + addressOf(c) + "] but was [" + realVal + "]";
         }
      }
      // Allow some time for all ClusteredGetCommands to finish executing
      TestingUtil.sleepThread(1000);
   }

   protected void assertOwnershipAndNonOwnership(Object key, boolean allowL1) {
      for (Cache<K, V> c : caches) {
         DataContainer dc = c.getAdvancedCache().getDataContainer();
         InternalCacheEntry ice = dc.get(key);
         if (isOwner(c, key)) {
            assert ice != null : "Fail on owner cache " + addressOf(c) + ": dc.get(" + key + ") returned null!";
            assert ice instanceof ImmortalCacheEntry : "Fail on owner cache " + addressOf(c) + ": dc.get(" + key + ") returned " + safeType(ice);
         } else {
            if (allowL1) {
               assert ice == null || ice instanceof MortalCacheEntry : "Fail on non-owner cache " + addressOf(c) + ": dc.get(" + key + ") returned " + safeType(ice);
            } else {
               assert ice == null : "Fail on non-owner cache " + addressOf(c) + ": dc.get(" + key + ") returned " + ice + "!";
            }
         }
      }
   }

   protected String safeType(Object o) {
      return DistributionTestHelper.safeType(o);
   }

   protected boolean isInL1(Cache<?, ?> cache, Object key) {
      DataContainer dc = cache.getAdvancedCache().getDataContainer();
      InternalCacheEntry ice = dc.get(key);
      return ice != null && !(ice instanceof ImmortalCacheEntry);
   }

   protected void assertIsInL1(Cache<?, ?> cache, Object key) {
      DistributionTestHelper.assertIsInL1(cache, key);
   }

   protected void assertIsNotInL1(Cache<?, ?> cache, Object key) {
      DistributionTestHelper.assertIsNotInL1(cache, key);
   }

   protected void assertIsInContainerImmortal(Cache<?, ?> cache, Object key) {
      DistributionTestHelper.assertIsInContainerImmortal(cache, key);
   }

   protected void assertIsInL1OrNull(Cache<?, ?> cache, Object key) {
      DistributionTestHelper.assertIsInL1OrNull(cache, key);
   }

   protected boolean isOwner(Cache<?, ?> c, Object key) {
      return DistributionTestHelper.isOwner(c, key);
   }

   protected boolean isFirstOwner(Cache<?, ?> c, Object key) {
      return DistributionTestHelper.isFirstOwner(c, key);
   }

   protected Cache<K, V>[] getOwners(Object key) {
      return getOwners(key, 2);
   }

   protected Cache<K, V>[] getOwners(Object key, int expectedNumberOwners) {
      Cache<K, V>[] owners = new Cache[expectedNumberOwners];
      int i = 0;
      for (Cache<K, V> c : caches) {
         if (isFirstOwner(c, key)) {
            owners[i++] = c;
            break;
         }
      }
      for (Cache<K, V> c : caches) {
         if (isOwner(c, key) && !isFirstOwner(c, key)) owners[i++] = c;
      }
      for (Cache<?, ?> c : owners) assert c != null : "Have not found enough owners for key [" + key + "]";
      return owners;
   }

   protected Cache<K, V>[] getNonOwnersExcludingSelf(Object key, Address self) {
      Cache<K, V>[] nonOwners = getNonOwners(key);
      boolean selfInArray = false;
      for (Cache<?, ?> c : nonOwners) {
         if (addressOf(c).equals(self)) {
            selfInArray = true;
            break;
         }
      }

      if (selfInArray) {
         Cache<K, V>[] nonOwnersExclSelf = new Cache[nonOwners.length - 1];
         int i = 0;
         for (Cache<K, V> c : nonOwners) {
            if (!addressOf(c).equals(self)) nonOwnersExclSelf[i++] = c;
         }
         return nonOwnersExclSelf;
      } else {
         return nonOwners;
      }
   }

   protected Cache<K, V>[] getNonOwners(Object key) {
      return getNonOwners(key, 2);
   }

   protected Cache<K, V>[] getNonOwners(Object key, int expectedNumberNonOwners) {
      Cache<K, V>[] nonOwners = new Cache[expectedNumberNonOwners];
      int i = 0;
      for (Cache<K, V> c : caches) {
         if (!isOwner(c, key)) nonOwners[i++] = c;
      }
      return nonOwners;
   }

   protected List<Address> residentAddresses(Object key) {
      DistributionManager dm = c1.getAdvancedCache().getComponentRegistry().getComponent(DistributionManager.class);
      return dm.locate(key);
   }

   protected DistributionManager getDistributionManager(Cache<?, ?> c) {
      return c.getAdvancedCache().getComponentRegistry().getComponent(DistributionManager.class);
   }

   protected ConsistentHash getConsistentHash(Cache<?, ?> c) {
      return getDistributionManager(c).getConsistentHash();
   }

   /**
    * Blocks and waits for a replication event on async caches
    *
    * @param key     key that causes the replication.  Used to determine which caches to listen on.  If null, all caches
    *                are checked
    * @param command command to listen for
    * @param caches  on which this key should be invalidated
    */
   protected void asyncWait(Object key, Class<? extends VisitableCommand> command, Cache<?, ?>... caches) {
      // no op.
   }

   protected TransactionManager getTransactionManager(Cache<?, ?> cache) {
      return TestingUtil.getTransactionManager(cache);
   }

   protected static void removeAllBlockingInterceptorsFromCache(Cache<?, ?> cache) {
      InterceptorChain chain = TestingUtil.extractComponent(cache, InterceptorChain.class);
      for (CommandInterceptor interceptor : chain.getInterceptorsWhichExtend(BlockingInterceptor.class)) {
         chain.removeInterceptor(interceptor.getClass());
      }
   }
}
