package com.atlassian.vcache.internal.guava;
   
   import com.atlassian.marshalling.api.MarshallingPair;
   import com.atlassian.vcache.CasIdentifier;
   import com.atlassian.vcache.DirectExternalCache;
   import com.atlassian.vcache.ExternalCacheException;
   import com.atlassian.vcache.IdentifiedValue;
   import com.atlassian.vcache.PutPolicy;
   import com.atlassian.vcache.internal.RequestContext;
  import com.atlassian.vcache.internal.core.ExternalCacheKeyGenerator;
  import com.atlassian.vcache.internal.core.cas.IdentifiedData;
  import com.atlassian.vcache.internal.core.service.AbstractExternalCache;
  import com.atlassian.vcache.internal.core.service.AbstractExternalCacheRequestContext;
  import com.atlassian.vcache.internal.core.service.UnversionedExternalCacheRequestContext;
  import com.google.common.cache.Cache;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  import java.util.HashMap;
  import java.util.Map;
  import java.util.Objects;
  import java.util.Optional;
  import java.util.Set;
  import java.util.concurrent.CompletionStage;
  import java.util.function.Function;
  import java.util.function.Supplier;
  import java.util.stream.Collectors;
  import java.util.stream.StreamSupport;
  
  import static com.atlassian.vcache.VCacheUtils.unsafeJoin;
  import static com.atlassian.vcache.internal.core.VCacheCoreUtils.isEmpty;
  import static com.atlassian.vcache.internal.core.cas.IdentifiedUtils.marshall;
  import static com.atlassian.vcache.internal.core.cas.IdentifiedUtils.safeCast;
  import static com.atlassian.vcache.internal.core.cas.IdentifiedUtils.unmarshall;
  import static com.atlassian.vcache.internal.core.cas.IdentifiedUtils.unmarshallIdentified;
  import static java.util.Objects.requireNonNull;
  
  /**
   * Guava based implementation of {@link DirectExternalCache}.
   *
   * @param <V> the value type.
   * @since 1.0.0
   */
  public class GuavaDirectExternalCache<V>
          extends AbstractExternalCache<V>
          implements DirectExternalCache<V> {
      private static final Logger log = LoggerFactory.getLogger(GuavaDirectExternalCache.class);
  
      private final Cache<String, IdentifiedData> delegate;
      private final Supplier<RequestContext> contextSupplier;
      private final ExternalCacheKeyGenerator keyGenerator;
      private final Optional<MarshallingPair<V>> valueMarshalling;
  
      public GuavaDirectExternalCache(
              String name,
              Cache<String, IdentifiedData> delegate,
              Supplier<RequestContext> contextSupplier,
              ExternalCacheKeyGenerator keyGenerator,
              Optional<MarshallingPair<V>> valueMarshalling) {
          super(name);
          this.delegate = requireNonNull(delegate);
          this.contextSupplier = requireNonNull(contextSupplier);
          this.keyGenerator = requireNonNull(keyGenerator);
          this.valueMarshalling = requireNonNull(valueMarshalling);
      }
  
      @Override
      public CompletionStage<Optional<V>> get(String internalKey) {
          return perform(() -> {
              final String externalKey = buildExternalKey(internalKey);
              final IdentifiedData identifiedData = delegate.getIfPresent(externalKey);
              return unmarshall(identifiedData, valueMarshalling);
          });
      }
  
      @Override
      public CompletionStage<V> get(String internalKey, Supplier<V> supplier) {
          return perform(() -> {
              final String externalKey = buildExternalKey(internalKey);
              final IdentifiedData identifiedData =
                      delegate.get(externalKey, () -> marshall(supplier.get(), valueMarshalling));
              return unmarshall(identifiedData, valueMarshalling).get();
          });
      }
  
      @Override
      public CompletionStage<Optional<IdentifiedValue<V>>> getIdentified(String internalKey) {
          return perform(() -> {
              final String externalKey = buildExternalKey(internalKey);
              return unmarshallIdentified(delegate.getIfPresent(externalKey), valueMarshalling);
          });
      }
  
      @Override
      public CompletionStage<IdentifiedValue<V>> getIdentified(String internalKey, Supplier<V> supplier) {
          return perform(() -> {
              final String externalKey = buildExternalKey(internalKey);
              return unmarshallIdentified(
                      delegate.get(externalKey, () -> marshall(supplier.get(), valueMarshalling)),
                     valueMarshalling).get();
         });
     }
 
     @Override
     public CompletionStage<Map<String, Optional<V>>> getBulk(Iterable<String> internalKeys) {
         return perform(() -> {
             if (isEmpty(internalKeys)) {
                 return new HashMap<>();
             }
 
             // De-duplicate the keys, calculate the externalKeys and retrieve
             final AbstractExternalCacheRequestContext cacheContext = ensureCacheContext();
             return StreamSupport.stream(internalKeys.spliterator(), false)
                     .distinct()
                     .collect(Collectors.toMap(
                             Objects::requireNonNull,
                             k -> unmarshall(delegate.getIfPresent(cacheContext.externalEntryKeyFor(k)), valueMarshalling)));
         });
     }
 
     @Override
     public CompletionStage<Map<String, V>> getBulk(
             Function<Set<String>, Map<String, V>> factory, Iterable<String> internalKeys) {
         return perform(() -> {
             if (isEmpty(internalKeys)) {
                 return new HashMap<>();
             }
 
             // De-duplicate the keys and then obtain the existing values
             final AbstractExternalCacheRequestContext cacheContext = ensureCacheContext();
             final Map<String, Optional<V>> existingValues = unsafeJoin(getBulk(internalKeys));
 
             // Add known values to the grand result
             final Map<String, V> grandResult = existingValues.entrySet().stream()
                     .filter(e -> e.getValue().isPresent())
                     .collect(Collectors.toMap(
                             Map.Entry::getKey,
                             e -> e.getValue().get()));
 
             // Bail out if we have all the values
             if (grandResult.size() == existingValues.size()) {
                 return grandResult;
             }
 
             // Sadly we now need to call the factory to create the missing values and then merge into the grand result.
             final Set<String> missingInternalKeys = existingValues.entrySet().stream()
                     .filter(e -> !e.getValue().isPresent())
                     .map(Map.Entry::getKey)
                     .collect(Collectors.toSet());
 
             final Map<String, V> missingValues = factory.apply(missingInternalKeys);
 
             missingValues.entrySet().forEach(e -> {
                 final Optional<V> existing = unmarshall(
                         delegate.asMap().putIfAbsent(
                                 cacheContext.externalEntryKeyFor(e.getKey()),
                                 marshall(e.getValue(), valueMarshalling)),
                         valueMarshalling);
                 grandResult.put(e.getKey(), existing.orElse(e.getValue()));
             });
 
             return grandResult;
         });
     }
 
     @Override
     public CompletionStage<Map<String, Optional<IdentifiedValue<V>>>> getBulkIdentified(Iterable<String> internalKeys) {
         return perform(() -> {
             if (isEmpty(internalKeys)) {
                 return new HashMap<>();
             }
 
             // De-duplicate the keys, calculate the externalKeys and retrieve
             final AbstractExternalCacheRequestContext cacheContext = ensureCacheContext();
             return StreamSupport.stream(internalKeys.spliterator(), false)
                     .distinct()
                     .collect(Collectors.toMap(
                             Objects::requireNonNull,
                             k -> unmarshallIdentified(delegate.getIfPresent(cacheContext.externalEntryKeyFor(k)), valueMarshalling)));
         });
     }
 
 
     @Override
     public CompletionStage<Boolean> put(String internalKey, V value, PutPolicy policy) {
         return perform(() -> {
             final AbstractExternalCacheRequestContext cacheContext = ensureCacheContext();
             final String externalKey = cacheContext.externalEntryKeyFor(internalKey);
             final IdentifiedData identifiedData = marshall(value, valueMarshalling);
 
             return GuavaUtils.directPut(externalKey, identifiedData, policy, delegate);
         });
     }
 
     @Override
     public CompletionStage<Boolean> removeIf(String internalKey, CasIdentifier casId) {
         return perform(() -> {
             final AbstractExternalCacheRequestContext cacheContext = ensureCacheContext();
             final String externalKey = cacheContext.externalEntryKeyFor(internalKey);
             final IdentifiedData existingData = safeCast(casId);
             return delegate.asMap().remove(externalKey, existingData);
         });
     }
 
     @Override
     public CompletionStage<Boolean> replaceIf(String internalKey, CasIdentifier casId, V newValue) {
         return perform(() -> {
             final AbstractExternalCacheRequestContext cacheContext = ensureCacheContext();
             final String externalKey = cacheContext.externalEntryKeyFor(internalKey);
             final IdentifiedData existingData = safeCast(casId);
             final IdentifiedData newData = marshall(newValue, valueMarshalling);
             return delegate.asMap().replace(externalKey, existingData, newData);
         });
     }
 
     @Override
     public CompletionStage<Void> remove(Iterable<String> internalKeys) {
         return perform(() -> {
             final AbstractExternalCacheRequestContext cacheContext = ensureCacheContext();
 
             StreamSupport.stream(internalKeys.spliterator(), false)
                     .distinct()
                     .map(cacheContext::externalEntryKeyFor)
                     .forEach(k -> delegate.asMap().remove(k));
             return null;
         });
     }
 
     @Override
     public CompletionStage<Void> removeAll() {
         return perform(() -> {
             delegate.asMap().clear();
             return null;
         });
     }
 
     protected AbstractExternalCacheRequestContext<V> ensureCacheContext() {
         final RequestContext requestContext = contextSupplier.get();
 
         return requestContext.computeIfAbsent(this, () -> {
             log.trace("Cache {}: Setting up a new context", getName());
             return new UnversionedExternalCacheRequestContext<>(
                     keyGenerator, getName(), requestContext::partitionIdentifier);
         });
     }
 
     @Override
     protected ExternalCacheException mapException(Exception ex) {
         return GuavaUtils.mapException(ex);
     }
 
     @Override
     protected Logger getLogger() {
         return log;
     }
 
     private String buildExternalKey(String internalKey) {
         final AbstractExternalCacheRequestContext cacheContext = ensureCacheContext();
         return cacheContext.externalEntryKeyFor(internalKey);
     }
 }