package com.atlassian.vcache.internal.test;
   
   import com.atlassian.vcache.internal.NameValidator;
   import com.atlassian.vcache.internal.RequestContext;
   import org.slf4j.Logger;
   import org.slf4j.LoggerFactory;
   
   import java.util.Map;
   import java.util.Optional;
  import java.util.concurrent.ConcurrentHashMap;
  import java.util.function.Supplier;
  
  import static com.atlassian.vcache.internal.NameValidator.requireValidPartitionIdentifier;
  import static java.util.Objects.requireNonNull;
  
  /**
   * Implementation of {@link Supplier} for {@link RequestContext} that manages an instance
   * per thread.
   * <p>
   * Notes:
   * </p>
   * <ul>
   * <li>
   * This class scopes thread-local variables to the instance. So two instances of {@link ThreadLocalRequestContextSupplier}
   * will have separate instances of {@link ThreadLocal}.
   * </li>
   * <li>
   * Using a lenient supplier <b>will break</b> {@link com.atlassian.vcache.TransactionalExternalCache}s. So be warned.
   * </li>
   * <li>
   * Before a thread can call {@link #get()}, the method {@link #initThread(String)} must
   * have been called, unless a lenient supplier is created.
   * </li>
   * <li>
   * The methods {@link #initThread(String)} and {@link #clearThread()} are not re-entrant.
   * Each call to {@link #initThread(String)} must be paired with a subsequent call to {@link #clearThread()}.
   * </li>
   * </ul>
   *
   * @since 1.0.0
   */
  public class ThreadLocalRequestContextSupplier implements Supplier<RequestContext> {
      private static final Logger log = LoggerFactory.getLogger(ThreadLocalRequestContextSupplier.class);
      private final ThreadLocal<RequestContext> threadRequestContexts = new ThreadLocal<>();
      private final Optional<Supplier<String>> lenientPartitionIdSupplier;
  
      private ThreadLocalRequestContextSupplier(Optional<Supplier<String>> lenientPartitionIdSupplier) {
          this.lenientPartitionIdSupplier = requireNonNull(lenientPartitionIdSupplier);
      }
  
      /**
       * Returns a strict supplier.
       */
      public static ThreadLocalRequestContextSupplier strictSupplier() {
          return new ThreadLocalRequestContextSupplier(Optional.empty());
      }
  
      /**
       * Returns a lenient supplier, which <b>will break</b> {@link com.atlassian.vcache.TransactionalExternalCache}s.
       *
       * @param partitionIdSupplier the supplier of partitionId's.
       */
      public static ThreadLocalRequestContextSupplier lenientSupplier(Supplier<String> partitionIdSupplier) {
          log.warn("A lenient supplier has been created, TransactionalExternalCaches are now broken");
          return new ThreadLocalRequestContextSupplier(Optional.of(partitionIdSupplier));
      }
  
      @Override
      public RequestContext get() {
          final RequestContext current = threadRequestContexts.get();
          if (current == null) {
              if (!lenientPartitionIdSupplier.isPresent()) {
                  log.error("Asked for request context when not initialised!");
                  throw new IllegalStateException("Thread has not been initialised.");
              }
  
              log.debug("Asked for request context when not initialised, returning a lenient one.");
              return new LenientRequestContext();
          }
  
          return current;
      }
  
      /**
       * Initialises the thread's {@link RequestContext}.
       *
       * @param partitionId the identifier for the partition. Will be validated using
       *                    {@link NameValidator#requireValidPartitionIdentifier(String)}.
       */
      public void initThread(String partitionId) {
          final RequestContext current = threadRequestContexts.get();
          if (current != null) {
              log.error(
                      "Asked to initialise thread {} that is already initialised!",
                      Thread.currentThread().getName());
              throw new IllegalStateException(
                      "Thread '" + Thread.currentThread().getName() + "' has already been initialised.");
          }
  
         requireValidPartitionIdentifier(partitionId);
         log.trace("Initialise request context");
         threadRequestContexts.set(new TestThreadLocalRequestContext(() -> partitionId));
     }
 
     public void clearThread() {
         final RequestContext current = threadRequestContexts.get();
         if (current == null) {
             log.warn("Asked to clear a thread that is already clear!");
         }
 
         log.trace("Clear request context");
         threadRequestContexts.remove();
     }
 
     /**
      * Implementation that knows and remembers nothing.
      */
     private class LenientRequestContext implements RequestContext {
         @Override
         public String partitionIdentifier() {
             return requireValidPartitionIdentifier(lenientPartitionIdSupplier.get().get());
         }
 
         @Override
         public <T> T computeIfAbsent(Object key, Supplier<T> supplier) {
             return supplier.get();
         }
 
         @Override
         public <T> Optional<T> get(Object key) {
             return Optional.empty();
         }
     }
 
     private static class TestThreadLocalRequestContext implements RequestContext {
 
         private final Supplier<String> partitionIdSupplier;
         private String partitionId;
         private final Map<Object, Object> map = new ConcurrentHashMap<>();
 
         public TestThreadLocalRequestContext(Supplier<String> partitionIdSupplier) {
             this.partitionIdSupplier = requireNonNull(partitionIdSupplier);
         }
 
         @Override
         public String partitionIdentifier() {
             // It doesn't matter if multiple threads get the value from the supplier, it should be a relatively cheap call
             // and it should always return the same value.
             if (partitionId == null) {
                 partitionId = partitionIdSupplier.get();
             }
             return partitionId;
         }
 
         @SuppressWarnings("unchecked")
         @Override
         public <T> T computeIfAbsent(Object key, Supplier<T> supplier) {
 
             return (T) map.computeIfAbsent(requireNonNull(key), o -> requireNonNull(supplier.get()));
         }
 
         @SuppressWarnings("unchecked")
         @Override
         public <T> Optional<T> get(Object key) {
             return Optional.ofNullable((T) map.get(key));
         }
     }
 }