package com.atlassian.vcache.internal.core.service;
   
   import com.atlassian.vcache.PutPolicy;
   import com.atlassian.vcache.internal.core.ExternalCacheKeyGenerator;
   import com.atlassian.vcache.internal.core.RecursionDetector;
   import com.google.common.collect.BiMap;
   import com.google.common.collect.HashBiMap;
   import com.google.common.collect.Maps;
   import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  import java.util.Collections;
  import java.util.HashMap;
  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.ConcurrentHashMap;
  import java.util.function.BiFunction;
  import java.util.function.Supplier;
  
  import static java.util.Objects.requireNonNull;
  
  /**
   * Represents the request context for an {@link com.atlassian.vcache.ExternalCache}.
   *
   * @param <V> the value type
   * @since 1.0.0
   */
  public abstract class AbstractExternalCacheRequestContext<V> {
      private static final Logger log = LoggerFactory.getLogger(AbstractExternalCacheRequestContext.class);
  
      protected final String name;
  
      /**
       * Used to track inflight calls to {@link #computeValue(String, BiFunction)} to detect recursive calls on
       * the same key. This is because these calls will deadlock.
       */
      private final RecursionDetector<String> recursionDetector = new RecursionDetector<>();
      private final ExternalCacheKeyGenerator keyGenerator;
      private final Supplier<String> partitionSupplier;
  
      // Need to synchronise the map to ensure consistency. Easier to use the Guava BiMap to achieve the result.
      private final BiMap<String, String> internalToExternalKeyMap = Maps.synchronizedBiMap(HashBiMap.create());
      // Make the inverse map unmodifiable, to ensure no duplicate processing.
      private final Map<String, String> externalToInternalKeyMap = Collections.unmodifiableMap(
              internalToExternalKeyMap.inverse());
  
      private final Map<String, CompletionStage<Optional<V>>> internalKeyToValueMap = new ConcurrentHashMap<>();
      private boolean hasRemoveAll;
      private final Map<String, DeferredOperation<V>> keyedOperationMap = new HashMap<>();
  
      protected AbstractExternalCacheRequestContext(ExternalCacheKeyGenerator keyGenerator,
                                                    String name,
                                                    Supplier<String> partitionSupplier) {
          this.keyGenerator = requireNonNull(keyGenerator);
          this.name = requireNonNull(name);
          this.partitionSupplier = requireNonNull(partitionSupplier);
      }
  
      protected abstract long cacheVersion();
  
      protected void clearKeyMaps() {
          internalToExternalKeyMap.clear();
      }
  
      public String externalEntryKeyFor(String internalKey) {
          final String cached = internalToExternalKeyMap.get(requireNonNull(internalKey));
          if (cached != null) {
              return cached;
          }
  
          final String result = keyGenerator.entryKey(
                  partitionSupplier.get(), name, cacheVersion(), internalKey);
  
          // Do a forcePut in case there is a race condition and another thread is also adding the entry
          internalToExternalKeyMap.forcePut(internalKey, result);
          return result;
      }
  
      /**
       * Returns the internal key for a supplied external key. The {@link #externalEntryKeyFor(String)} must have
       * returned the value now being specified, in the current request.
       *
       * @param externalKey the external key previously returned by {@link #externalEntryKeyFor(String)}
       * @return the internal key
       */
      public String internalEntryKeyFor(String externalKey) {
          return requireNonNull(externalToInternalKeyMap.get(externalKey));
      }
  
      public Optional<Optional<V>> getValueRecorded(String internalKey) {
          final CompletionStage<Optional<V>> optionalCompletionStage = internalKeyToValueMap.get(requireNonNull(internalKey));
          if (optionalCompletionStage == null || optionalCompletionStage.toCompletableFuture().isCompletedExceptionally()) {
              return Optional.empty();
          }
          return Optional.ofNullable(optionalCompletionStage.toCompletableFuture().join());
      }
 
     public void recordValue(String internalKey, Optional<V> outcome) {
         recordValue(requireNonNull(internalKey), CompletableFuture.completedFuture(outcome));
     }
 
     public void recordValue(String internalKey, CompletionStage<Optional<V>> outcome) {
         log.trace("Cache {}, recording value for {}", name, internalKey);
         internalKeyToValueMap.put(requireNonNull(internalKey), outcome);
     }
 
     public void recordValues(Map<String, V> knownValues) {
         log.trace("Cache {}, recording {} known values", name, knownValues.size());
         knownValues.entrySet().stream().collect(
                 () -> internalKeyToValueMap,
                 (m, e) -> m.put(e.getKey(), CompletableFuture.completedFuture(Optional.of(e.getValue()))),
                 Map::putAll);
     }
 
     public CompletionStage<Optional<V>> computeValue(
             String key,
             BiFunction<String, CompletionStage<Optional<V>>, CompletionStage<Optional<V>>> doCompute) {
         try (RecursionDetector.Guard ignored = recursionDetector.guardOn(key)) {
             return internalKeyToValueMap.compute(requireNonNull(key), doCompute);
         }
     }
 
     public void forgetValue(String internalKey) {
         log.trace("Cache {}, forgetting value for {}", name, internalKey);
         internalKeyToValueMap.remove(internalKey);
     }
 
     public void forgetAllValues() {
         log.trace("Cache {}, forgetting all values", name);
         internalKeyToValueMap.clear();
     }
 
     /**
      * Saves a new value in the value map and adds the put policy for this operation.
      *
      * @param internalKey The internal key.
      * @param value       The value to record.
      * @param policy      The {@link PutPolicy} used to record this value.
      */
     public void recordPut(String internalKey, V value, PutPolicy policy) {
         recordValue(internalKey, Optional.of(value));
         recordPutPolicy(internalKey, value, policy);
     }
 
     /**
      * Stores the policy for the recordPut().
      *
      * @param internalKey The internal key
      * @param value       The value saved.
      * @param policy      The {@link PutPolicy} used to record this value.
      */
     public void recordPutPolicy(String internalKey, V value, PutPolicy policy) {
         keyedOperationMap.put(requireNonNull(internalKey), DeferredOperation.putOperation(value, policy));
     }
 
     public void recordRemove(Iterable<String> internalKeys) {
         for (String internalKey : internalKeys) {
             recordValue(internalKey, Optional.empty());
             keyedOperationMap.put(requireNonNull(internalKey), DeferredOperation.removeOperation());
         }
     }
 
     public void recordRemoveAll() {
         forgetAllValues();
         hasRemoveAll = true;
         keyedOperationMap.clear(); // Forget everything, as starting again
     }
 
     public boolean hasRemoveAll() {
         return hasRemoveAll;
     }
 
     public void forgetAll() {
         forgetAllValues();
         hasRemoveAll = false;
         keyedOperationMap.clear();
     }
 
     public Set<Map.Entry<String, DeferredOperation<V>>> getKeyedOperations() {
         return keyedOperationMap.entrySet();
     }
 
     public boolean hasPendingOperations() {
         return hasRemoveAll || !keyedOperationMap.isEmpty();
     }
 
     /**
      * Represents a deferred operation.
      *
      * @param <V> the value type
      */
     public static class DeferredOperation<V> {
         private final boolean remove;
         private final CompletionStage<Optional<V>> value;
         private final Optional<PutPolicy> policy;
 
         private DeferredOperation() {
             this.remove = true;
             this.value = CompletableFuture.completedFuture(Optional.empty());
             this.policy = Optional.empty();
         }
 
         private DeferredOperation(V value, PutPolicy policy) {
             this(CompletableFuture.completedFuture(Optional.of(value)), policy);
         }
 
         private DeferredOperation(CompletionStage<Optional<V>> value, PutPolicy policy) {
             this.remove = false;
             this.value = requireNonNull(value);
             this.policy = Optional.of(policy);
         }
 
         public boolean isRemove() {
             return remove;
         }
 
         public boolean isPut() {
             return !remove;
         }
 
         public V getValue() {
             return value.toCompletableFuture().join().get();
         }
 
         public PutPolicy getPolicy() {
             return policy.get();
         }
 
         public static <V> DeferredOperation<V> removeOperation() {
             return new DeferredOperation<>();
         }
 
         public static <V> DeferredOperation<V> putOperation(V value, PutPolicy policy) {
             return new DeferredOperation<>(value, policy);
         }
     }
 }