package example.memory;

import java.io.IOException;
import java.lang.management.BufferPoolMXBean;
import java.lang.management.ManagementFactory;
import java.lang.management.MemoryUsage;
import java.lang.ref.Reference;
import java.nio.ByteBuffer;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.nio.file.Path;
import java.time.Duration;
import java.time.Instant;
import java.util.ArrayList;
import java.util.List;

public final class MemoryRegionProbe {
    private static final int MIB = 1024 * 1024;
    private static final int HEAP_MIB = 16;
    private static final int DIRECT_MIB = 12;
    private static final Duration WAIT_LIMIT = Duration.ofMinutes(5);
    private static volatile List<byte[]> retainedHeap;
    private static volatile ByteBuffer retainedDirect;

    private MemoryRegionProbe() {
    }

    public static void main(String[] args) throws Exception {
        if (args.length == 1 && "direct-limit".equals(args[0])) {
            proveDirectLimit();
            return;
        }
        if (args.length != 1) {
            throw new IllegalArgumentException("expected a control directory or direct-limit");
        }
        runControlled(Path.of(args[0]));
    }

    private static void runControlled(Path control) throws Exception {
        Files.createDirectories(control);
        BufferPoolMXBean directPool = directPool();
        MemoryUsage heapBefore = ManagementFactory.getMemoryMXBean().getHeapMemoryUsage();
        long directCountBefore = directPool.getCount();
        long directCapacityBefore = directPool.getTotalCapacity();
        long directUsedBefore = directPool.getMemoryUsed();

        write(control.resolve("pid"), Long.toString(ProcessHandle.current().pid()));
        write(control.resolve("ready"), "ready");
        await(control.resolve("allocate"));

        List<byte[]> heap = new ArrayList<>(HEAP_MIB);
        long heapPayload = 0L;
        for (int i = 0; i < HEAP_MIB; i++) {
            byte[] block = new byte[MIB];
            touch(block, (byte) (i + 1));
            heap.add(block);
            heapPayload += block.length;
        }

        ByteBuffer direct = ByteBuffer.allocateDirect(DIRECT_MIB * MIB);
        for (int i = 0; i < direct.capacity(); i += 4096) {
            direct.put(i, (byte) 0x5a);
        }
        direct.put(direct.capacity() - 1, (byte) 0x5a);
        retainedHeap = heap;
        retainedDirect = direct;

        MemoryUsage heapAfter = ManagementFactory.getMemoryMXBean().getHeapMemoryUsage();
        long directCountDelta = directPool.getCount() - directCountBefore;
        long directCapacityDelta = directPool.getTotalCapacity() - directCapacityBefore;
        long directUsedAfter = directPool.getMemoryUsed();
        long directUsedDelta = directUsedBefore < 0L || directUsedAfter < 0L
                ? -1L
                : directUsedAfter - directUsedBefore;
        long heapDelta = heapAfter.getUsed() - heapBefore.getUsed();

        boolean heapPayloadLive = retainedHeap.size() == HEAP_MIB
                && heapPayload == (long) HEAP_MIB * MIB
                && retainedHeap.get(HEAP_MIB - 1)[0] == (byte) HEAP_MIB;
        boolean directPayloadLive = retainedDirect.isDirect()
                && retainedDirect.capacity() == DIRECT_MIB * MIB
                && retainedDirect.get(retainedDirect.capacity() - 1) == (byte) 0x5a;
        boolean directPoolObserved = directCountDelta >= 1L
                && directCapacityDelta >= (long) DIRECT_MIB * MIB;
        if (!heapPayloadLive || !directPayloadLive || !directPoolObserved) {
            throw new AssertionError("memory boundary assertions failed");
        }

        String report = String.join(System.lineSeparator(),
                "heap-requested-bytes=" + heapPayload,
                "heap-used-delta-observed=" + heapDelta,
                "heap-payload-live=" + heapPayloadLive,
                "direct-requested-bytes=" + direct.capacity(),
                "direct-count-delta=" + directCountDelta,
                "direct-capacity-delta=" + directCapacityDelta,
                "direct-memory-used-delta-observed=" + directUsedDelta,
                "direct-payload-live=" + directPayloadLive,
                "direct-pool-observed=" + directPoolObserved)
                + System.lineSeparator();
        write(control.resolve("report"), report);
        write(control.resolve("allocated"), "allocated");
        await(control.resolve("stop"));
        Reference.reachabilityFence(heap);
        Reference.reachabilityFence(direct);
    }

    private static void proveDirectLimit() {
        MemoryUsage heap = ManagementFactory.getMemoryMXBean().getHeapMemoryUsage();
        long heapHeadroom = heap.getMax() - heap.getUsed();
        boolean heapHasHeadroom = heap.getMax() > 0L && heapHeadroom > 8L * MIB;
        if (!heapHasHeadroom) {
            throw new AssertionError("heap headroom must exceed the direct allocation request");
        }
        try {
            ByteBuffer.allocateDirect(8 * MIB);
        } catch (OutOfMemoryError expected) {
            String message = expected.getMessage();
            boolean directBufferFailure = message != null && message.contains("direct buffer memory");
            if (!directBufferFailure) {
                throw new AssertionError("unexpected OutOfMemoryError kind", expected);
            }
            System.out.println("direct-limit=direct-buffer-OOME+heap-headroom");
            return;
        }
        throw new AssertionError("8 MiB direct allocation unexpectedly passed the 4 MiB limit");
    }

    private static BufferPoolMXBean directPool() {
        return ManagementFactory.getPlatformMXBeans(BufferPoolMXBean.class).stream()
                .filter(pool -> "direct".equals(pool.getName()))
                .findFirst()
                .orElseThrow(() -> new IllegalStateException("direct buffer pool MXBean is unavailable"));
    }

    private static void touch(byte[] block, byte marker) {
        for (int i = 0; i < block.length; i += 4096) {
            block[i] = marker;
        }
        block[block.length - 1] = marker;
    }

    private static void await(Path signal) throws InterruptedException {
        Instant deadline = Instant.now().plus(WAIT_LIMIT);
        while (!Files.exists(signal)) {
            if (Instant.now().isAfter(deadline)) {
                throw new IllegalStateException("timed out waiting for " + signal.getFileName());
            }
            Thread.sleep(25L);
        }
    }

    private static void write(Path path, String value) throws IOException {
        Files.writeString(path, value, StandardCharsets.UTF_8);
    }
}