package net.openhft.chronicle.core.util;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.lang.ref.Reference;
import java.lang.ref.ReferenceQueue;
import java.lang.ref.WeakReference;
import java.util.Collections;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.ConcurrentLinkedQueue;
import java.util.concurrent.Executor;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicIntegerFieldUpdater;
import java.util.function.Supplier;

/**
 * Intended as a replacement for usage of sun.misc.Cleaner/jdk.internal.ref.Cleaner
 */
public final class WeakReferenceCleaner extends WeakReference<Object> {
    private static final Logger LOGGER = LoggerFactory.getLogger(WeakReferenceCleaner.class);
    private static final ReferenceQueue<Object> REFERENCE_QUEUE = new ReferenceQueue<>();
    private static final ConcurrentLinkedQueue<WeakReferenceCleaner> SCHEDULED_CLEAN = new ConcurrentLinkedQueue<>();
    private static final Map<WeakReferenceCleaner, Boolean> REFERENCE_MAP =
            Collections.synchronizedMap(
                    new HashMap<>(128));
    private static final AtomicBoolean REFERENCE_PROCESSOR_STARTED = new AtomicBoolean(false);
    private static final AtomicIntegerFieldUpdater<WeakReferenceCleaner> CLEANED_FLAG =
            AtomicIntegerFieldUpdater.newUpdater(WeakReferenceCleaner.class, "cleaned");

    private final Runnable thunk;
    @SuppressWarnings("unused")
    private volatile int cleaned = 0;

    private WeakReferenceCleaner(final Object referent, final Runnable thunk) {
        super(referent, REFERENCE_QUEUE);
        this.thunk = thunk;
    }

    public static WeakReferenceCleaner newCleaner(final Object referent, final Runnable thunk) {
        startReferenceProcessor(WeakReferenceCleaner::referenceCleanerExecutor);

        final WeakReferenceCleaner cleaner = new WeakReferenceCleaner(referent, thunk);
        REFERENCE_MAP.put(cleaner, Boolean.TRUE);
        return cleaner;
    }

    public static void startReferenceProcessor(final Supplier<Executor> executorSupplier) {
        if (!REFERENCE_PROCESSOR_STARTED.get()) {
            if (REFERENCE_PROCESSOR_STARTED.compareAndSet(false, true)) {
                executorSupplier.get().execute(new ReferenceProcessor());
            }
        }
    }

    static Executor referenceCleanerExecutor() {
        ExecutorService executor = Executors.newSingleThreadExecutor(r -> {
            final Thread t = new Thread(r);
            t.setName("chronicle-weak-reference-cleaner");
            t.setDaemon(true);
            return t;
        });

        Runtime.getRuntime().addShutdownHook(new Thread(executor::shutdown));

        return executor;
    }

    public static int referenceCount() {
        return REFERENCE_MAP.size();
    }

    public void clean() {
        if (CLEANED_FLAG.compareAndSet(this, 0, 1)) {
            thunk.run();
        }
    }

    public void scheduleForClean() {
        SCHEDULED_CLEAN.add(this);
        REFERENCE_MAP.remove(this);
    }

    static final class ReferenceProcessor implements Runnable {
        @Override
        public void run() {
            Thread thread = Thread.currentThread();
            while (!thread.isInterrupted()) {
                try {
                    // prioritise scheduled cleaners
                    WeakReferenceCleaner wrc;
                    while ((wrc = SCHEDULED_CLEAN.poll()) != null) {
                        wrc.clean();
                    }

                    Reference<?> reference;
                    while ((reference = REFERENCE_QUEUE.remove(50L)) != null) {
                        final WeakReferenceCleaner cleaner = (WeakReferenceCleaner) reference;

                        cleaner.clean();
                    }

                } catch (InterruptedException e) {
                    LOGGER.debug("Interrupted while trying to retrieve reference, exiting.", e);
                    thread.interrupt();
                    return;
                } catch (Throwable e) {
                    LOGGER.warn("Exception while trying to process reference.", e);
                }
            }
        }
    }
}

@@ -219,67 +219,9 @@
 til.
-Collections;%0Aimport java.util.HashMap;%0Aimport java.util.Map
+*
 ;%0Aim
@@ -998,32 +998,211 @@
 inkedQueue%3C%3E();%0A
+    /*%0A     * This set is used to hold a STRONG reference to the WeakReferenceCleaner object, to avoid it being collected%0A     * before GC added it to the reference queue%0A     */%0A
     private stat
@@ -1210,19 +1210,19 @@
 c final 
-Map
+Set
 %3CWeakRef
@@ -1234,25 +1234,16 @@
 eCleaner
-, Boolean
 %3E REFERE
@@ -1250,25 +1250,13 @@
 NCE_
-MAP =%0A           
+SET =
  Col
@@ -1280,41 +1280,20 @@
 ized
-Map(%0A                    
+Set(
 new Hash
 Map%3C
@@ -1292,11 +1292,11 @@
 Hash
-Map
+Set
 %3C%3E(1
@@ -2125,23 +2125,23 @@
 NCE_
-MAP.put
+SET.add
 (cleaner
 , Bo
@@ -2140,22 +2140,8 @@
 aner
-, Boolean.TRUE
 );%0A 
@@ -2312,28 +2312,12 @@
 et()
-) %7B%0A            if (
+ && 
 REFE
@@ -2367,18 +2367,16 @@
 , true))
- %7B
 %0A       
@@ -2446,32 +2446,8 @@
 ));%0A
-            %7D%0A        %7D%0A
     
@@ -2858,93 +2858,8 @@
  %7D%0A%0A
-    public static int referenceCount() %7B%0A        return REFERENCE_MAP.size();%0A    %7D%0A%0A
     
@@ -3070,11 +3070,11 @@
 NCE_
-MAP
+SET
 .rem
@@ -3768,16 +3768,70 @@
 erence;%0A
+                        REFERENCE_SET.remove(cleaner);
 %0A       

Fix reference map not being cleaned up properly