package com.atlassian.vcache.internal.core.service;
   
   import com.atlassian.vcache.ExternalCacheException;
   import com.atlassian.vcache.PutPolicy;
   import com.atlassian.vcache.TransactionalExternalCache;
   import com.atlassian.vcache.internal.MetricLabel;
   import com.atlassian.vcache.internal.RequestContext;
   import com.atlassian.vcache.internal.core.TransactionControl;
   import com.atlassian.vcache.internal.core.metrics.CacheType;
  import com.atlassian.vcache.internal.core.metrics.MetricsRecorder;
  
  import java.util.Collections;
  import java.util.HashMap;
  import java.util.Map;
  import java.util.Optional;
  import java.util.Set;
  import java.util.concurrent.ExecutionException;
  import java.util.function.Function;
  import java.util.function.Supplier;
  import java.util.stream.Collectors;
  
  import static java.util.Objects.requireNonNull;
  
  /**
   * Provides operations common for {@link com.atlassian.vcache.TransactionalExternalCache} instances.
   *
   * @param <V> the value type
   * @since 1.0.0
   */
  public abstract class AbstractTransactionalExternalCache<V>
          extends AbstractNonDirectExternalCache<V>
          implements TransactionalExternalCache<V>, TransactionControl {
  
      protected final Supplier<RequestContext> contextSupplier;
  
      protected AbstractTransactionalExternalCache(String name, Supplier<RequestContext> contextSupplier, MetricsRecorder metricsRecorder) {
          super(name, metricsRecorder);
          this.contextSupplier = requireNonNull(contextSupplier);
      }
  
      @Override
      public final void put(String internalKey, V value, PutPolicy policy) {
          final AbstractExternalCacheRequestContext<V> cacheContext = ensureCacheContext();
          cacheContext.recordPut(internalKey, value, policy);
      }
  
      @Override
      public final void remove(Iterable<String> internalKeys) {
          final AbstractExternalCacheRequestContext<V> cacheContext = ensureCacheContext();
          cacheContext.recordRemove(internalKeys);
      }
  
      @Override
      public final void removeAll() {
          final AbstractExternalCacheRequestContext<V> cacheContext = ensureCacheContext();
          cacheContext.recordRemoveAll();
      }
  
      @Override
      public final boolean transactionDiscard() {
          final RequestContext requestContext = contextSupplier.get();
          final Optional<AbstractExternalCacheRequestContext<V>> cacheRequestContext = requestContext.get(this);
  
          if (!cacheRequestContext.isPresent()) {
              // there are no pending operations
              return false;
          }
  
          final boolean hasPendingOperations = cacheRequestContext.get().hasPendingOperations();
          cacheRequestContext.get().forgetAll();
          return hasPendingOperations;
      }
  
      protected final Map<String, Optional<V>> checkValuesRecorded(Iterable<String> internalKeys) {
          final AbstractExternalCacheRequestContext<V> cacheContext = ensureCacheContext();
  
          final Map<String, Optional<V>> result = new HashMap<>();
  
          internalKeys.forEach(k -> {
              final Optional<Optional<V>> valueRecorded = cacheContext.getValueRecorded(k);
              if (valueRecorded.isPresent()) {
                  result.put(k, valueRecorded.get());
              } else if (cacheContext.hasRemoveAll()) {
                  result.put(k, Optional.empty());
              }
          });
  
          return result;
      }
  
      @Override
      protected final V handleCreation(String internalKey, Supplier<V> supplier) throws ExecutionException, InterruptedException {
          final AbstractExternalCacheRequestContext<V> cacheContext = ensureCacheContext();
          final V suppliedValue = requireNonNull(supplier.get());
          cacheContext.recordPutPolicy(internalKey, suppliedValue, PutPolicy.ADD_ONLY);
          return suppliedValue;
      }
  
      @Override
     protected final 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.
         final AbstractExternalCacheRequestContext<V> cacheContext = ensureCacheContext();
 
         // Need to handle if removeAll has been performed, and hence not check remotely. Otherwise,
         // need to check remotely.
         final Map<String, Optional<V>> candidateValues;
         if (cacheContext.hasRemoveAll()) {
             candidateValues = new HashMap<>();
         } else {
             metricsRecorder.record(name, CacheType.EXTERNAL, MetricLabel.NUMBER_OF_REMOTE_GET, 1);
             candidateValues = directGetBulk(externalKeys);
         }
 
         final Set<String> missingExternalKeys = cacheContext.hasRemoveAll() ?
                 externalKeys :
                 candidateValues.entrySet().stream()
                         .filter(e -> !e.getValue().isPresent())
                         .map(Map.Entry::getKey)
                         .collect(Collectors.toSet());
 
         // Add the retrieved values to the grand result
         final Map<String, V> grandResult = candidateValues.entrySet().stream()
                 .filter(e -> e.getValue().isPresent())
                 .collect(Collectors.toMap(
                         e -> cacheContext.internalEntryKeyFor(e.getKey()),
                         e -> e.getValue().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 record adding them
             missingValues.entrySet().forEach(e -> put(e.getKey(), e.getValue(), PutPolicy.ADD_ONLY));
 
             grandResult.putAll(missingValues);
         }
 
         return grandResult;
     }
 }