package com.atlassian.vcache.internal.memcached;
   
   import com.atlassian.marshalling.api.MarshallingPair;
   import com.atlassian.vcache.ExternalCacheException;
   import com.atlassian.vcache.ExternalCacheSettings;
   import com.atlassian.vcache.PutPolicy;
   import com.atlassian.vcache.internal.MetricLabel;
   import com.atlassian.vcache.internal.RequestContext;
   import com.atlassian.vcache.internal.core.ExternalCacheKeyGenerator;
  import com.atlassian.vcache.internal.core.VCacheCoreUtils;
  import com.atlassian.vcache.internal.core.metrics.CacheType;
  import com.atlassian.vcache.internal.core.metrics.MetricsRecorder;
  import com.atlassian.vcache.internal.core.service.AbstractStableReadExternalCache;
  import com.atlassian.vcache.internal.core.service.VersionedExternalCacheRequestContext;
  import net.spy.memcached.MemcachedClientIF;
  import net.spy.memcached.OperationTimeoutException;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  import java.util.Collections;
  import java.util.HashSet;
  import java.util.Map;
  import java.util.Optional;
  import java.util.Set;
  import java.util.concurrent.CompletableFuture;
  import java.util.concurrent.CompletionStage;
  import java.util.concurrent.ExecutionException;
  import java.util.concurrent.Future;
  import java.util.concurrent.atomic.AtomicBoolean;
  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.internal.core.VCacheCoreUtils.isEmpty;
  import static com.atlassian.vcache.internal.core.VCacheCoreUtils.marshall;
  import static com.atlassian.vcache.internal.core.VCacheCoreUtils.unmarshall;
  import static com.atlassian.vcache.internal.memcached.MemcachedUtils.expiryTime;
  import static com.atlassian.vcache.internal.memcached.MemcachedUtils.putOperationForPolicy;
  import static java.util.Objects.requireNonNull;
  
  /**
   * Implementation that backs onto Memcached.
   *
   * @param <V> the value type
   * @since 1.0.0
   */
  class MemcachedStableReadExternalCache<V>
          extends AbstractStableReadExternalCache<V> {
      private static final Logger log = LoggerFactory.getLogger(MemcachedStableReadExternalCache.class);
  
      private final Supplier<MemcachedClientIF> clientSupplier;
      private final Supplier<RequestContext> contextSupplier;
      private final ExternalCacheKeyGenerator keyGenerator;
      private final MarshallingPair<V> valueMarshalling;
      private final int ttlSeconds;
  
      MemcachedStableReadExternalCache(
              Supplier<MemcachedClientIF> clientSupplier,
              Supplier<RequestContext> contextSupplier,
              ExternalCacheKeyGenerator keyGenerator,
              String name,
              MarshallingPair<V> valueMarshalling,
              ExternalCacheSettings settings,
              MetricsRecorder metricsRecorder) {
          super(name, metricsRecorder);
          this.clientSupplier = requireNonNull(clientSupplier);
          this.contextSupplier = requireNonNull(contextSupplier);
          this.keyGenerator = requireNonNull(keyGenerator);
          this.valueMarshalling = requireNonNull(valueMarshalling);
          this.ttlSeconds = VCacheCoreUtils.roundUpToSeconds(settings.getDefaultTtl().get());
      }
  
      @Override
      public CompletionStage<Boolean> internalPut(String internalKey, V value, PutPolicy policy) {
          final VersionedExternalCacheRequestContext<V> cacheContext = ensureCacheContext();
          return perform(() -> {
                      // This stores the result of the remote operation. If it fails then we can roll back. We use Atomic because in the
                      // future we may be async.
                      final AtomicBoolean remotePutResult = new AtomicBoolean(true);
                      cacheContext.computeValue(internalKey, (key, oldValue) -> {
  
                          // computeValue allows us to get the current CompletableFuture (if there is one) and chain the next
                          // put() onto the old future. If the old future is still running then we preserve
                          // ordering (and so consistency). Also any blocked gets will get the new value when the chain
                          // completes. The map locks the update to the future so updates are nicely serialized.
                          if (oldValue != null) {
                              return oldValue.thenApply((oldV) -> {
                                  remotePutResult.set(remotePut(cacheContext.externalEntryKeyFor(internalKey), value, policy));
                                  return Optional.of(value);
                              });
                          } else {
                              remotePutResult.set(remotePut(cacheContext.externalEntryKeyFor(internalKey), value, policy));
                              return CompletableFuture.completedFuture(Optional.of(value));
                          }
                      });
                      return remotePutResult.get();
                  },
                  (result) -> {
                     if (!result) {
                         cacheContext.forgetValue(internalKey);
                     }
                 });
     }
 
     private boolean remotePut(String externalKey, V value, PutPolicy policy) {
 
         final byte[] valueBytes = valueMarshalling.getMarshaller().marshallToBytes(value);
 
         try {
             final Future<Boolean> putOp = putOperationForPolicy(policy, clientSupplier.get(), externalKey, expiryTime(ttlSeconds), valueBytes);
             return putOp.get();
         } catch (InterruptedException | ExecutionException e) {
             throw new ExternalCacheException(ExternalCacheException.Reason.UNCLASSIFIED_FAILURE, e);
         }
     }
 
     @Override
     protected CompletionStage<Void> internalRemove(Iterable<String> internalKeys) {
         // There is no bulk delete in the api, so need to remove each one async
         return perform(() -> {
             if (isEmpty(internalKeys)) {
                 return null;
             }
 
             // Lodge all the requests for delete
             final VersionedExternalCacheRequestContext<V> cacheContext = ensureCacheContext();
             final Map<String, Future<Boolean>> deleteOps =
                     StreamSupport.stream(internalKeys.spliterator(), false)
                             .distinct()
                             .collect(Collectors.toMap(
                                     k -> k,
                                     k -> clientSupplier.get().delete(cacheContext.externalEntryKeyFor(k))
                             ));
 
             // Need to loop verifying the outcome. If an exception was thrown for an operation,
             // then do not record the removal.
             Exception failureException = null;
             for (Map.Entry<String, Future<Boolean>> delOp : deleteOps.entrySet()) {
                 try {
                     delOp.getValue().get();
                     cacheContext.recordValue(delOp.getKey(), Optional.empty());
                 } catch (ExecutionException | InterruptedException ex) {
                     log.info("Cache {}: unable to remove key {}", name, delOp.getKey(), ex);
                     failureException = ex;
                 }
             }
 
             // Finally, if there was a failure, then re-throw the last one reported (as good as any!)
             if (failureException != null) {
                 if (failureException instanceof ExecutionException) {
                     throw (ExecutionException) failureException;
                 }
                 throw (InterruptedException) failureException;
             }
 
             return null;
         });
     }
 
     @Override
     protected CompletionStage<Void> internalRemoveAll() {
         return perform(() -> {
             final VersionedExternalCacheRequestContext<V> cacheContext = ensureCacheContext();
             cacheContext.updateCacheVersion(
                     MemcachedUtils.incrementCacheVersion(clientSupplier, cacheContext.externalCacheVersionKey()));
             cacheContext.forgetAllValues();
             return null;
         });
     }
 
     @Override
     protected Logger getLogger() {
         return log;
     }
 
     @Override
     protected VersionedExternalCacheRequestContext<V> ensureCacheContext() throws OperationTimeoutException {
         final RequestContext requestContext = contextSupplier.get();
 
         return requestContext.computeIfAbsent(this, () -> {
             // Need to build a new context, which involves getting the current cache version, or setting it if it does
             // not exist.
             log.trace("Cache {}: Setting up a new context", name);
             final VersionedExternalCacheRequestContext<V> newCacheContext = new VersionedExternalCacheRequestContext<>(
                     keyGenerator, name, requestContext::partitionIdentifier);
             newCacheContext.updateCacheVersion(
                     MemcachedUtils.obtainCacheVersion(clientSupplier, newCacheContext.externalCacheVersionKey()));
             return newCacheContext;
         });
     }
 
     @Override
     protected V handleCreation(String internalKey, Supplier<V> supplier)
             throws ExecutionException, InterruptedException {
         final VersionedExternalCacheRequestContext<V> cacheContext = ensureCacheContext();
         final V candidateValue = requireNonNull(supplier.get());
         final byte[] candidateValueBytes = valueMarshalling.getMarshaller().marshallToBytes(candidateValue);
         final String externalKey = cacheContext.externalEntryKeyFor(internalKey);
 
         // Loop until either able to add the candidate value, or retrieve one that has been added by another thread
         for (; ; ) {
             final Future<Boolean> addOp = clientSupplier.get().add(externalKey, expiryTime(ttlSeconds), candidateValueBytes);
             if (addOp.get()) {
                 // I break here, rather than just return, due to battling with the compiler. Unless written
                 // this way, the compiler will not allow the lambda structure.
                 break;
             }
 
             getLogger().info("Cache {}, unable to add candidate for key {}, retrieve what was added", name, internalKey);
             final Optional<V> otherAddedValue = unmarshall((byte[]) clientSupplier.get().get(externalKey), valueMarshalling);
             if (otherAddedValue.isPresent()) {
                 return otherAddedValue.get();
             }
 
             getLogger().info("Cache {}, unable to retrieve recently added candidate for key {}, looping", name, internalKey);
         }
 
         return candidateValue;
     }
 
     @Override
     protected Map<String, V> handleCreation(Function<Set<String>, Map<String, V>> factory, Set<String> externalKeys)
             throws ExecutionException, InterruptedException {
         // Get the missing values from the external cache.
         // Returns map of keys that contain values, so need to handle the missing ones
         final VersionedExternalCacheRequestContext<V> cacheContext = ensureCacheContext();
         metricsRecorder.record(name, CacheType.EXTERNAL, MetricLabel.NUMBER_OF_REMOTE_GET, 1);
         final Map<String, Object> haveValues = clientSupplier.get().getBulk(externalKeys);
         getLogger().trace("Cache {}: getBulk(Function): {} of {} entries have values",
                 name, haveValues.size(), externalKeys.size());
         final Set<String> missingExternalKeys = new HashSet<>(externalKeys);
         missingExternalKeys.removeAll(haveValues.keySet());
 
         // Add the retrieved values to the grand result
         final Map<String, V> grandResult = haveValues.entrySet().stream()
                 .collect(Collectors.toMap(
                         e -> cacheContext.internalEntryKeyFor(e.getKey()),
                         e -> unmarshall((byte[]) e.getValue(), valueMarshalling).get()));
 
         if (!missingExternalKeys.isEmpty()) {
             getLogger().trace("Cache {}: getBulk(Function): calling factory to create {} values",
                     name, missingExternalKeys.size());
             // Okay, need to get the missing values and mapping from externalKeys to internalKeys
             final Set<String> missingInternalKeys = Collections.unmodifiableSet(
                     missingExternalKeys.stream().map(cacheContext::internalEntryKeyFor).collect(Collectors.toSet()));
             final Map<String, V> missingValues = factory.apply(missingInternalKeys);
             if (missingInternalKeys.size() != missingValues.size()) {
                 getLogger().warn("Cache {}: getBulk(Function): mismatch on generated values, expected ",
                         name, missingInternalKeys.size() + " but got " + missingValues.size());
                 throw new ExternalCacheException(ExternalCacheException.Reason.FUNCTION_INCORRECT_RESULT);
             }
 
             // Okay, got the missing values, now need to add them to Memcached
             final Map<String, Future<Boolean>> internalKeyToFutureMap = missingValues.entrySet().stream().collect(Collectors.toMap(
                     Map.Entry::getKey,
                     e -> clientSupplier.get().set(
                             cacheContext.externalEntryKeyFor(e.getKey()), expiryTime(ttlSeconds), marshall(e.getValue(), valueMarshalling))
             ));
 
             // Now wait for the outcomes and then add to the grand result
             for (Map.Entry<String, Future<Boolean>> e : internalKeyToFutureMap.entrySet()) {
                 e.getValue().get(); // Don't care about the result as it will always be true
             }
 
             grandResult.putAll(missingValues);
         }
 
         return grandResult;
     }
 
     @Override
     protected final ExternalCacheException mapException(Exception ex) {
         return MemcachedUtils.mapException(ex);
     }
 
     @Override
     protected final Optional<V> directGet(String externalKey) {
         return unmarshall((byte[]) clientSupplier.get().get(externalKey), valueMarshalling);
     }
 
     @Override
     protected final Map<String, Optional<V>> directGetBulk(Set<String> externalKeys) {
         return MemcachedUtils.directGetBulk(externalKeys, clientSupplier, valueMarshalling);
     }
 }