package org.infinispan.distribution.rehash;

import static org.testng.Assert.assertEquals;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Random;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.Future;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;

import javax.transaction.Transaction;
import javax.transaction.TransactionManager;
import javax.transaction.xa.XAResource;
import javax.transaction.xa.Xid;

import org.infinispan.Cache;
import org.infinispan.distribution.BaseDistFunctionalTest;
import org.infinispan.distribution.MagicKey;
import org.infinispan.test.TestingUtil;
import org.testng.annotations.Test;

/**
 * A base test for all rehashing tests
 */
@Test(groups = "functional")
public abstract class RehashTestBase extends BaseDistFunctionalTest<Object, String> {

   protected RehashTestBase() {
      cleanup = CleanupPhase.AFTER_METHOD;
      tx = true;
      performRehashing = true;
   }

   // this setup has 4 running caches: {c1, c2, c3, c4}

   /**
    * This is overridden by subclasses.  Could typically be a JOIN or LEAVE event.
    * @param offline
    */
   abstract void performRehashEvent(boolean offline) throws Throwable;

   /**
    * Blocks until a rehash completes.
    */
   abstract void waitForRehashCompletion();

   void additionalWait() {
      TestingUtil.sleepThread(1000);
   }

   protected List<MagicKey> init() {

      List<MagicKey> keys = new ArrayList<>(Arrays.asList(
            new MagicKey("k1", c1), new MagicKey("k2", c2),
            new MagicKey("k3", c3), new MagicKey("k4", c4)
      ));
      assertEquals(caches.size(), keys.size(), "Received caches" + caches);

      int i = 0;
      for (Cache<Object, String> c : caches) c.put(keys.get(i++), "v" + i);

      i = 0;
      for (MagicKey key : keys) assertOwnershipAndNonOwnership(key, false);

      log.infof("Initialized with keys %s", keys);
      return keys;
   }

   /**
    * Simple test.  Put some state, trigger event, test results
    */
   @Test
   public void testNonTransactional() throws Throwable {
      List<MagicKey> keys = init();

      log.info("Invoking rehash event");
      performRehashEvent(false);

      waitForRehashCompletion();
      log.info("Rehash complete");

      int i = 0;
      for (MagicKey key : keys) assertOnAllCachesAndOwnership(key, "v" + ++i);
   }


   /**
    * More complex - init some state.  Start a new transaction, and midway trigger a rehash.  Then complete transaction
    * and test results.
    */
   @Test
   public void testTransactional() throws Throwable {
      final List<MagicKey> keys = init();
      final CountDownLatch l = new CountDownLatch(1);
      final AtomicBoolean rollback = new AtomicBoolean(false);

      Future<Void> future = fork(new Runnable() {
         @Override
         public void run() {
            try {
               // start a transaction on c1.
               TransactionManager t1 = TestingUtil.getTransactionManager(c1);
               t1.begin();
               c1.put(keys.get(0), "transactionally_replaced");
               Transaction tx = t1.getTransaction();
               tx.enlistResource(new XAResourceAdapter() {
                  public int prepare(Xid id) {
                     // this would be called *after* the cache prepares.
                     try {
                        log.debug("Unblocking commit");
                        l.await();
                     } catch (InterruptedException e) {
                        Thread.currentThread().interrupt();
                     }
                     return XAResource.XA_OK;
                  }
               });
               t1.commit();
            } catch (Exception e) {
               log.error("Error committing transaction", e);
               rollback.set(true);
               throw new RuntimeException(e);
            }
         }
      }, null);

      log.info("Invoking rehash event");
      performRehashEvent(true);
      l.countDown();
      future.get(30, TimeUnit.SECONDS);

      //ownership can only be verified after the rehashing has completed
      waitForRehashCompletion();
      log.info("Rehash complete");

      //only check for these values if tx was not rolled back
      if (!rollback.get()) {
         // the ownership of k1 might change during the tx and a cache might end up with it in L1
         assertOwnershipAndNonOwnership(keys.get(0), true);
         assertOwnershipAndNonOwnership(keys.get(1), false);
         assertOwnershipAndNonOwnership(keys.get(2), false);
         assertOwnershipAndNonOwnership(keys.get(3), false);

         // checking the values will bring the keys to L1, so we want to do it after checking ownership
         assertOnAllCaches(keys.get(0), "transactionally_replaced");
         assertOnAllCaches(keys.get(1), "v" + 2);
         assertOnAllCaches(keys.get(2), "v" + 3);
         assertOnAllCaches(keys.get(3), "v" + 4);
      }
   }

   /**
    * A stress test.  One node is constantly modified while a rehash occurs.
    */
   @Test(groups = "stress")
   public void testNonTransactionalStress() throws Throwable {
      stressTest(false);
   }

   /**
    * A stress test.  One node is constantly modified using transactions while a rehash occurs.
    */
   @Test(groups = "stress")
   public void testTransactionalStress() throws Throwable {
      stressTest(true);
   }

   private void stressTest(boolean tx) throws Throwable {
      final List<MagicKey> keys = init();
      final CountDownLatch latch = new CountDownLatch(1);
      List<Updater> updaters = new ArrayList<>(keys.size());
      for (MagicKey k : keys) {
         Updater u = new Updater(c1, k, latch, tx);
         u.start();
         updaters.add(u);
      }

      latch.countDown();

      log.info("Invoking rehash event");
      performRehashEvent(false);

      for (Updater u : updaters) u.complete();
      for (Updater u : updaters) u.join();

      waitForRehashCompletion();

      log.info("Rehash complete");

      int i = 0;
      for (MagicKey key : keys) assertOnAllCachesAndOwnership(key, "v" + updaters.get(i++).currentValue);
   }
}

class Updater extends Thread {
   static final Random r = new Random();
   volatile int currentValue = 0;
   MagicKey key;
   Cache cache;
   CountDownLatch latch;
   volatile boolean running = true;
   TransactionManager tm;

   Updater(Cache cache, MagicKey key, CountDownLatch latch, boolean tx) {
      super("Updater-" + key);
      this.key = key;
      this.cache = cache;
      this.latch = latch;
      if (tx) tm = TestingUtil.getTransactionManager(cache);
   }

   public void complete() {
      running = false;
   }

   @Override
   public void run() {
      while (running) {
         try {
            currentValue++;
            if (tm != null) tm.begin();
            cache.put(key, "v" + currentValue);
            if (tm != null) tm.commit();
            TestingUtil.sleepThread(r.nextInt(10) * 10);
         } catch (Exception e) {
            // do nothing?
         }
      }
   }
}
