package com.atlassian.cache.ehcache.cluster;
   
   import java.util.List;
   import java.util.concurrent.atomic.AtomicInteger;
   
   import com.atlassian.cache.Cache;
   import com.atlassian.cache.CachedReference;
   import com.atlassian.cache.ManagedCache;
   import com.atlassian.cache.ehcache.EhCacheManager;
  
  import com.google.common.base.Function;
  import com.google.common.base.Throwables;
  import com.google.common.collect.ImmutableList;
  import com.google.common.collect.Iterables;
  
  import org.junit.rules.ExternalResource;
  
  import net.sf.ehcache.CacheManager;
  import net.sf.ehcache.config.Configuration;
  
  import static com.atlassian.cache.ehcache.cluster.EhCacheConfigurationFactory.newConfiguration;
  
  /**
   * JUnit rule to set up an Ehcache cluster for testing
   */
  public class EhCacheCluster extends ExternalResource
  {
      private final AtomicInteger activeCount = new AtomicInteger();
      private final int count;
  
      private volatile List<CacheManager> rawNodes = ImmutableList.of();
      private volatile List<EhCacheManager> decoratedNodes = ImmutableList.of();
  
      public EhCacheCluster(final int count)
      {
          this.count = count;
      }
  
      @Override
      protected void after()
      {
          try
          {
              if (activeCount.decrementAndGet() == 0)
              {
                  // only truly shut down when the active count reaches zero
                  final List<CacheManager> nodes = rawNodes;
                  decoratedNodes = null;
                  rawNodes = null;
                  for (CacheManager instance : nodes)
                  {
                      instance.shutdown();
                  }
              }
          }
          finally
          {
              super.after();
          }
      }
  
      @Override
      protected void before() throws Throwable
      {
          super.before();
  
          if (activeCount.getAndIncrement() == 0)
          {
              // only initialize when the active count is at zero
              rawNodes = EhCacheCluster.newCluster(count);
              decoratedNodes = ImmutableList.copyOf(Iterables.transform(rawNodes,
                      new Function<CacheManager, EhCacheManager>()
                      {
                          @Override
                          public EhCacheManager apply(CacheManager rawNode)
                          {
                              return new EhCacheManager(rawNode, null);
                          }
                      }));
          }
      }
  
      public EhCacheManager getNode(int index)
      {
          return decoratedNodes.get(index);
      }
  
      public void reset()
      {
          if (size() > 0)
          {
              final EhCacheManager node = decoratedNodes.get(0);
              for (ManagedCache cache : node.getManagedCaches())
              {
                  if (cache instanceof Cache<?,?>)
                  {
                      ((Cache<?,?>)cache).removeAll();
                  }
                  else
                 {
                     ((CachedReference<?>)cache).reset();
                 }
             }
         }
     }
 
     public int size()
     {
         return decoratedNodes.size();
     }
 
     public static EhCacheCluster createCluster(int count)
     {
         return new EhCacheCluster(count);
     }
 
     public static List<CacheManager> newCluster(final int maxNodeId)
     {
         final ImmutableList.Builder<CacheManager> nodes = new ImmutableList.Builder<CacheManager>();
         for (int thisNodeId = 1; thisNodeId <= maxNodeId; ++thisNodeId)
         {
             final Configuration config = newConfiguration(thisNodeId, maxNodeId);
             final CacheManager ehcache = CacheManager.newInstance(config);
             nodes.add(ehcache);
         }
         return nodes.build();
     }
 
     public static void clusteredTest(int nodeCount, ClusteredTest test) throws Exception
     {
         final EhCacheCluster cluster = new EhCacheCluster(nodeCount);
         try
         {
             cluster.before();
             try
             {
                 test.runTest(cluster);
             }
             finally
             {
                 cluster.after();
             }
         }
         catch (Throwable e)
         {
             Throwables.propagate(e);
         }
     }
 
     public interface ClusteredTest
     {
         void runTest(EhCacheCluster cluster) throws Exception;
     }
 }