package example.roots;

import com.sun.management.HotSpotDiagnosticMXBean;
import java.lang.management.ManagementFactory;
import java.lang.ref.Reference;
import java.lang.ref.ReferenceQueue;
import java.lang.ref.WeakReference;
import java.util.ArrayList;
import java.util.List;

public final class AllocationRootsSafepointProbe {
    private static final int PAYLOAD_BYTES = 4 * 1024 * 1024;
    private static final int SMALL_ALLOCATION_COUNT = 10_000;
    private static final int SMALL_ALLOCATION_BYTES = 128;
    private static final int RELEASE_ATTEMPTS = 20;
    private static final long QUEUE_WAIT_MILLIS = 100L;
    private static final List<byte[]> ROOT = new ArrayList<>();
    private static volatile long pressureChecksum;

    private AllocationRootsSafepointProbe() {
    }

    public static void main(String[] args) throws InterruptedException {
        if (args.length != 1) {
            throw new IllegalArgumentException(
                    "usage: AllocationRootsSafepointProbe <root-retained|root-released|allocation>");
        }
        switch (args[0]) {
            case "root-retained" -> rootRetained();
            case "root-released" -> rootReleased();
            case "allocation" -> allocation();
            default -> throw new IllegalArgumentException("unknown mode: " + args[0]);
        }
    }

    private static void rootRetained() {
        Observation observation = installPayload();
        System.gc();

        byte[] rooted = ROOT.get(0);
        boolean weakRefersToRoot = observation.reference().refersTo(rooted);
        boolean queueEmpty = observation.queue().poll() == null;
        require(ROOT.size() == 1, "the strong root disappeared");
        require(weakRefersToRoot, "the weak reference was cleared while a strong root remained");
        require(queueEmpty, "the weak reference was enqueued while a strong root remained");

        System.out.println("mode=root-retained");
        System.out.println("root-size=1");
        System.out.println("weak-refers-to-root=true");
        System.out.println("queue-empty=true");
        Reference.reachabilityFence(rooted);
        Reference.reachabilityFence(observation.reference());
    }

    private static void rootReleased() throws InterruptedException {
        Observation observation = installPayload();
        ROOT.clear();

        Reference<? extends byte[]> enqueued = null;
        for (int attempt = 0; attempt < RELEASE_ATTEMPTS && enqueued == null; attempt++) {
            allocationPressure();
            System.gc();
            enqueued = observation.queue().remove(QUEUE_WAIT_MILLIS);
        }

        require(ROOT.isEmpty(), "the strong root was not cleared");
        require(enqueued == observation.reference(), "the expected weak reference was not enqueued");
        require(observation.reference().get() == null, "the weak referent was not cleared");

        System.out.println("mode=root-released");
        System.out.println("root-size=0");
        System.out.println("observer-enqueued=true");
        System.out.println("weak-cleared=true");
        Reference.reachabilityFence(observation.reference());
    }

    private static void allocation() {
        long checksum = 0L;
        for (int index = 0; index < SMALL_ALLOCATION_COUNT; index++) {
            byte[] value = new byte[SMALL_ALLOCATION_BYTES];
            int slot = index & (SMALL_ALLOCATION_BYTES - 1);
            value[slot] = (byte) (index * 31);
            checksum += Byte.toUnsignedInt(value[slot]);
        }
        System.gc();

        System.out.println("mode=allocation");
        System.out.println("use-tlab=" + vmOption("UseTLAB"));
        System.out.println("allocation-count=" + SMALL_ALLOCATION_COUNT);
        System.out.println("allocation-bytes=" + SMALL_ALLOCATION_BYTES);
        System.out.println("checksum=" + checksum);
    }

    private static Observation installPayload() {
        ReferenceQueue<byte[]> queue = new ReferenceQueue<>();
        byte[] payload = new byte[PAYLOAD_BYTES];
        payload[0] = 42;
        payload[payload.length - 1] = 24;
        WeakReference<byte[]> reference = new WeakReference<>(payload, queue);
        ROOT.add(payload);
        return new Observation(reference, queue);
    }

    private static void allocationPressure() {
        long checksum = 0L;
        for (int index = 0; index < 512; index++) {
            byte[] value = new byte[4096];
            value[index & (value.length - 1)] = (byte) index;
            checksum += Byte.toUnsignedInt(value[index & (value.length - 1)]);
        }
        pressureChecksum = checksum;
    }

    private static String vmOption(String name) {
        HotSpotDiagnosticMXBean bean =
                ManagementFactory.getPlatformMXBean(HotSpotDiagnosticMXBean.class);
        require(bean != null, "HotSpotDiagnosticMXBean is unavailable");
        return bean.getVMOption(name).getValue();
    }

    private static void require(boolean condition, String message) {
        if (!condition) {
            throw new IllegalStateException(message);
        }
    }

    private record Observation(
            WeakReference<byte[]> reference,
            ReferenceQueue<byte[]> queue) {
    }
}