package com.atlassian.cache.ehcache.wrapper;
   
   import net.sf.ehcache.CacheException;
   import net.sf.ehcache.Ehcache;
   import net.sf.ehcache.Status;
   import net.sf.ehcache.loader.CacheLoader;
   
   import java.util.Collection;
   import java.util.HashMap;
  import java.util.Map;
  
  import static com.atlassian.cache.ehcache.wrapper.WrapperUtils.unwrapAllKeys;
  
  /**
   * This class makes sure that during any loading operation, wrapped keys do not leak to the end user.
   * If the cache that uses {@link ValueProcessor} decides to load the value, this decorator should be used on the provided loader.
   */
  public class ValueProcessorEhcacheLoaderDecorator implements CacheLoader {
  
      private final CacheLoader delegate;
      private final ValueProcessor valueProcessor;
  
      public ValueProcessorEhcacheLoaderDecorator(final CacheLoader delegate, final ValueProcessor valueProcessor) {
          this.delegate = delegate;
          this.valueProcessor = valueProcessor;
      }
  
      @Override
      public Object load(final Object key) throws CacheException {
          return valueProcessor.wrap(delegate.load(valueProcessor.unwrap(key)));
      }
  
      @SuppressWarnings("unchecked")
      @Override
      public Map loadAll(final Collection keys) {
          return wrapLoadedValues(delegate.loadAll(unwrapAllKeys(keys, valueProcessor)));
      }
  
      @Override
      public Object load(final Object key, final Object argument) {
          return valueProcessor.wrap(delegate.load(valueProcessor.unwrap(key), argument));
      }
  
      @SuppressWarnings("unchecked")
      @Override
      public Map loadAll(final Collection keys, final Object argument) {
          return wrapLoadedValues(delegate.loadAll(unwrapAllKeys(keys, valueProcessor), argument));
      }
  
      @Override
      public String getName() {
          return delegate.getName();
      }
  
      @Override
      public CacheLoader clone(final Ehcache cache) throws CloneNotSupportedException {
          return delegate.clone(cache);
      }
  
      @Override
      public void init() {
          delegate.init();
      }
  
      @Override
      public void dispose() throws CacheException {
          delegate.dispose();
      }
  
      @Override
      public Status getStatus() {
          return delegate.getStatus();
      }
  
      private Map<Object, Object> wrapLoadedValues(final Map<Object, Object> unwrappedValues) {
          Map<Object, Object> result = new HashMap<>(unwrappedValues.size());
  
          for (final Map.Entry<Object, Object> entry : unwrappedValues.entrySet()) {
              result.put(valueProcessor.wrap(entry.getKey()), valueProcessor.wrap(entry.getValue()));
          }
          return result;
      }
  }