package com.atlassian.vcache.internal.core.metrics;
   
   import com.atlassian.marshalling.api.MarshallingPair;
   import com.atlassian.vcache.DirectExternalCache;
   import com.atlassian.vcache.JvmCache;
   import com.atlassian.vcache.RequestCache;
   import com.atlassian.vcache.StableReadExternalCache;
   import com.atlassian.vcache.TransactionalExternalCache;
   import com.atlassian.vcache.internal.MetricLabel;
  import com.atlassian.vcache.internal.RequestContext;
  import com.atlassian.vcache.internal.RequestMetrics;
  import com.atlassian.vcache.internal.core.TransactionControl;
  
  import java.util.function.Supplier;
  
  import static java.util.Objects.requireNonNull;
  
  /**
   * Implementation of {@link MetricsCollector}.
   *
   * @since 1.0.0
   */
  public class DefaultMetricsCollector implements MetricsCollector {
  
      private final Supplier<RequestContext> contextSupplier;
  
      public DefaultMetricsCollector(Supplier<RequestContext> contextSupplier) {
          this.contextSupplier = requireNonNull(contextSupplier);
      }
  
      @Override
      public void record(String cacheName, CacheType cacheType, MetricLabel metricLabel, long sample) {
          obtainMetrics(contextSupplier.get()).record(cacheName, cacheType, metricLabel, sample);
      }
  
      @Override
      public RequestMetrics obtainRequestMetrics(RequestContext context) {
          return obtainMetrics(context);
      }
  
      @Override
      public TransactionControl wrap(TransactionControl control, String cacheName) {
          return new TimedTransactionControl(control, this, cacheName);
      }
  
      @Override
      public <T> MarshallingPair<T> wrap(MarshallingPair<T> marshalling, String cacheName) {
          return new MarshallingPair<>(
                  new TimedMarshaller<>(marshalling.getMarshaller(), this, cacheName),
                  new TimedUnmarshaller<>(marshalling.getUnmarshaller(), this, cacheName));
      }
  
      @Override
      public <K, V> JvmCache<K, V> wrap(JvmCache<K, V> cache) {
          return new TimedJvmCache<>(cache, this);
      }
  
      @Override
      public <K, V> RequestCache<K, V> wrap(RequestCache<K, V> cache) {
          return new TimedRequestCache<>(cache, this);
      }
  
      @Override
      public <V> DirectExternalCache<V> wrap(DirectExternalCache<V> cache) {
          return new TimedDirectExternalCache<>(this, cache);
      }
  
      @Override
      public <V> StableReadExternalCache<V> wrap(StableReadExternalCache<V> cache) {
          return new TimedStableReadExternalCache<>(this, cache);
      }
  
      @Override
      public <V> TransactionalExternalCache<V> wrap(TransactionalExternalCache<V> cache) {
          return new TimedTransactionalExternalCache<>(this, cache);
      }
  
      MutableRequestMetrics obtainMetrics(RequestContext context) {
          return context.computeIfAbsent(this, DefaultRequestMetrics::new);
      }
  }