/*
    * Licensed to the Apache Software Foundation (ASF) under one
    * or more contributor license agreements.  See the NOTICE file
    * distributed with this work for additional information
    * regarding copyright ownership.  The ASF licenses this file
    * to you under the Apache License, Version 2.0 (the
    * "License"); you may not use this file except in compliance
    * with the License.  You may obtain a copy of the License at
    *
   *   http://www.apache.org/licenses/LICENSE-2.0
   *
   * Unless required by applicable law or agreed to in writing,
   * software distributed under the License is distributed on an
   * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
   * KIND, either express or implied.  See the License for the
   * specific language governing permissions and limitations
   * under the License.
   */
  package org.eclipse.aether.util.concurrency;
  
  import java.lang.ref.ReferenceQueue;
  import java.lang.ref.WeakReference;
  import java.util.concurrent.ConcurrentHashMap;
  
  /**
   * A concurrent cache with weak keys and weak values, inspired by maven-impl's Cache class
   * but stripped down and optimized for hot-path performance.
   * <p>
   * Design compared to the full Cache:
   * <ul>
   *   <li><b>Zero allocation on get()</b> — uses a ThreadLocal reusable lookup key
   *       instead of allocating a new wrapper on every read</li>
   *   <li><b>O(1) cleanup per stale entry</b> — uses identity-based removal from
   *       the ReferenceQueue instead of a full {@code entrySet().removeIf()} scan</li>
   *   <li><b>Lock-free reads</b> — {@link ConcurrentHashMap#get} is a volatile read
   *       with no lock acquisition</li>
   *   <li><b>Lock-striped writes</b> — ConcurrentHashMap uses fine-grained locking</li>
   * </ul>
   * <p>
   * Keys are held via {@link WeakReference}: when a key is no longer strongly referenced
   * elsewhere, its entry becomes eligible for garbage collection. Values are also held
   * via WeakReference. Stale entries are cleaned up lazily on {@link #put}.
   *
   * @param <K> the type of keys
   * @param <V> the type of values
   * @since 2.0.19
   */
  public class ConcurrentWeakCache<K, V> {
  
      private final ConcurrentHashMap<Object, WeakReference<V>> map;
      private final ReferenceQueue<Object> queue = new ReferenceQueue<>();
  
      @SuppressWarnings("rawtypes")
      private static final ThreadLocal<LookupKey> LOOKUP_KEY = new ThreadLocal<LookupKey>() {
          @Override
          protected LookupKey initialValue() {
              return new LookupKey();
          }
      };
  
      /**
       * Creates a new cache with default initial capacity.
       */
      public ConcurrentWeakCache() {
          this.map = new ConcurrentHashMap<>();
      }
  
      /**
       * Creates a new cache with the given initial capacity.
       *
       * @param initialCapacity the initial capacity
       */
      public ConcurrentWeakCache(int initialCapacity) {
          this.map = new ConcurrentHashMap<>(initialCapacity);
      }
  
      /**
       * Returns the value for the given key, or {@code null} if the key is not present
       * or the value has been garbage collected.
       * <p>
       * This method is lock-free and allocates no objects.
       *
       * @param key the key to look up
       * @return the cached value, or {@code null}
       */
      @SuppressWarnings("unchecked")
      public V get(K key) {
          LookupKey<K> lookupKey = LOOKUP_KEY.get();
          lookupKey.set(key);
          try {
              WeakReference<V> ref = map.get(lookupKey);
              return ref != null ? ref.get() : null;
          } finally {
              lookupKey.clear();
          }
      }
  
      /**
       * Stores a key-value pair in the cache. Both key and value are held via weak references.
      * <p>
      * Also performs lazy cleanup of entries whose keys have been garbage collected.
      *
      * @param key   the key
      * @param value the value
      */
     public void put(K key, V value) {
         expungeStaleEntries();
         map.put(new WeakKey<>(key, queue), new WeakReference<>(value));
     }
 
     /**
      * If the key is not already present (or its value has been GC'd), stores the key-value pair
      * and returns the given value. If the key is already present with a live value, returns the
      * existing value without storing. Concurrent callers for the same key are guaranteed to
      * receive the same instance (the insert-or-keep decision is atomic per key via
      * {@link ConcurrentHashMap#merge}).
      * <p>
      * Also performs lazy cleanup of entries whose keys have been garbage collected.
      *
      * @param key   the key
      * @param value the value to store if absent
      * @return the existing value if present, or the given value if newly stored
      */
     public V putIfAbsent(K key, V value) {
         // Fast path: lock-free lookup, zero allocation (ThreadLocal LookupKey)
         V existing = get(key);
         if (existing != null) {
             return existing;
         }
         // Slow path: atomic insert-or-keep via merge() — locks only the target bucket,
         // so contention is per-key, not global. Handles the case where an existing entry's
         // value has been GC'd by atomically replacing it with the new value.
         expungeStaleEntries();
         WeakReference<V> newRef = new WeakReference<>(value);
         WeakReference<V> ref = map.merge(new WeakKey<>(key, queue), newRef, (existingRef, incoming) -> {
             V existingValue = existingRef.get();
             return existingValue != null ? existingRef : incoming;
         });
         if (ref != newRef) {
             V existingValue = ref.get();
             if (existingValue != null) {
                 return existingValue;
             }
         }
         return value;
     }
 
     /**
      * Returns the number of entries in the cache (including possibly stale ones).
      *
      * @return the cache size
      */
     public int size() {
         return map.size();
     }
 
     /**
      * Removes entries whose weak-reference keys have been garbage collected.
      * Called automatically on {@link #put}. Each stale entry is removed in O(1)
      * via identity match — no map scan required.
      */
     private void expungeStaleEntries() {
         Object staleKey;
         while ((staleKey = queue.poll()) != null) {
             // ConcurrentHashMap.remove() checks (storedKey == staleKey) first.
             // Since the staleKey IS the same WeakKey object that was stored in
             // the map, identity matches and the entry is removed in O(1).
             map.remove(staleKey);
         }
     }
 
     /**
      * Reusable lookup key stored in a ThreadLocal for zero-allocation reads.
      * Used only for {@link ConcurrentHashMap#get} lookups, never stored in the map.
      * <p>
      * Implements equals/hashCode to match {@link WeakKey} entries in the map:
      * ConcurrentHashMap.get(lookupKey) calls {@code lookupKey.equals(storedWeakKey)},
      * which delegates to {@code key.equals(storedWeakKey.get())}.
      */
     static class LookupKey<K> {
         private K key;
         private int hash;
 
         void set(K key) {
             this.key = key;
             this.hash = key.hashCode();
         }
 
         void clear() {
             this.key = null; // prevent leak via ThreadLocal
         }
 
         @Override
         public int hashCode() {
             return hash;
         }
 
         @Override
         @SuppressWarnings("rawtypes")
         public boolean equals(Object o) {
             if (o instanceof WeakKey) {
                 Object otherKey = ((WeakKey) o).get();
                 return key != null && key.equals(otherKey);
             }
             return false;
         }
     }
 
     /**
      * Weak-reference key stored in the map. When the referent is garbage collected,
      * this object is enqueued in the {@link ReferenceQueue} for cleanup.
      * <p>
      * The cleanup path uses identity: {@code map.remove(staleWeakKey)} matches because
      * ConcurrentHashMap checks {@code storedKey == staleWeakKey} before calling equals().
      * Since the queued WeakKey IS the same object that's in the map, this is always true.
      */
     static class WeakKey<K> extends WeakReference<K> {
         private final int hash;
 
         @SuppressWarnings("unchecked")
         WeakKey(K key, ReferenceQueue<? super K> queue) {
             super(key, (ReferenceQueue<? super K>) queue);
             this.hash = key.hashCode();
         }
 
         @Override
         public int hashCode() {
             return hash;
         }
 
         @Override
         @SuppressWarnings("rawtypes")
         public boolean equals(Object o) {
             if (this == o) {
                 return true; // identity match — used for cleanup and normal dedup
             }
             K k = get();
             if (k == null) {
                 return false; // referent GC'd
             }
             if (o instanceof WeakKey) {
                 return k.equals(((WeakKey) o).get());
             }
             if (o instanceof LookupKey) {
                 return k.equals(((LookupKey) o).key);
             }
             return false;
         }
     }
 }