package example.jvm.jit;

import java.io.IOException;
import java.io.InputStream;
import java.lang.invoke.MethodHandle;
import java.lang.invoke.MethodHandles;
import java.lang.invoke.MethodType;
import java.nio.file.Files;
import java.nio.file.Path;
import java.util.ArrayList;
import java.util.List;

public final class JitLifecycleProbe {
    private interface Operation {
        int apply(int value);
    }

    private static final class AddOne implements Operation {
        @Override
        public int apply(int value) {
            return value + 1;
        }
    }

    private static final class DoubleValue implements Operation {
        @Override
        public int apply(int value) {
            return value * 2;
        }
    }

    private static final class Point {
        private final int x;
        private final int y;

        private Point(int x, int y) {
            this.x = x;
            this.y = y;
        }
    }

    private static final class HotMethodTemplate {
        private static int hot(int value) {
            int mixed = value * 31 + 17;
            return (mixed ^ (mixed >>> 7)) + 3;
        }
    }

    private JitLifecycleProbe() {
    }

    private static int dispatch(Operation operation, int value) {
        return operation.apply(value);
    }

    private static long monomorphicLoop(int iterations) {
        Operation addOne = new AddOne();
        long checksum = 0;
        for (int i = 0; i < iterations; i++) {
            checksum += dispatch(addOne, i & 1023);
        }
        return checksum;
    }

    private static long bimorphicLoop(int iterations) {
        Operation addOne = new AddOne();
        Operation doubleValue = new DoubleValue();
        long checksum = 0;
        for (int i = 0; i < iterations; i++) {
            Operation operation = (i & 1) == 0 ? addOne : doubleValue;
            checksum += dispatch(operation, i & 1023);
        }
        return checksum;
    }

    private static long distance(int value) {
        Point point = new Point(value, value + 1);
        return (long) point.x * point.x + (long) point.y * point.y;
    }

    private static long allocationLoop(int iterations) {
        long checksum = 0;
        for (int i = 0; i < iterations; i++) {
            checksum += distance(i & 1023);
        }
        return checksum;
    }

    private static long semanticLoop(int iterations) {
        return monomorphicLoop(iterations)
                + bimorphicLoop(iterations)
                + allocationLoop(iterations);
    }

    private static byte[] hotMethodTemplateBytes() throws IOException {
        String resource = "/" + HotMethodTemplate.class.getName().replace('.', '/') + ".class";
        try (InputStream input = JitLifecycleProbe.class.getResourceAsStream(resource)) {
            if (input == null) {
                throw new IOException("missing class resource: " + resource);
            }
            return input.readAllBytes();
        }
    }

    private static long codeCachePressure(int classCount, int callsPerClass) throws Throwable {
        byte[] template = hotMethodTemplateBytes();
        MethodHandles.Lookup lookup = MethodHandles.lookup();
        MethodType methodType = MethodType.methodType(int.class, int.class);
        List<MethodHandle> handles = new ArrayList<>(classCount);
        long checksum = 0;
        for (int classIndex = 0; classIndex < classCount; classIndex++) {
            MethodHandles.Lookup hidden = lookup.defineHiddenClass(template, true);
            MethodHandle handle = hidden.findStatic(hidden.lookupClass(), "hot", methodType);
            handles.add(handle);
            for (int call = 0; call < callsPerClass; call++) {
                checksum += (int) handle.invokeExact((call + classIndex) & 1023);
            }
        }
        if (handles.size() != classCount) {
            throw new IllegalStateException("hidden classes were not retained");
        }
        return checksum;
    }

    private static int parsePositive(String[] args, int index, int fallback) {
        if (args.length <= index) {
            return fallback;
        }
        int value = Integer.parseInt(args[index]);
        if (value <= 0) {
            throw new IllegalArgumentException("iteration or hold value must be positive");
        }
        return value;
    }

    public static void main(String[] args) throws Throwable {
        String mode = args.length == 0 ? "semantic" : args[0];
        switch (mode) {
            case "semantic" -> {
                int iterations = parsePositive(args, 1, 500_000);
                System.out.println("mode=semantic checksum=" + semanticLoop(iterations));
            }
            case "lifecycle" -> {
                int iterations = parsePositive(args, 1, 5_000_000);
                System.out.println("phase=monomorphic");
                long monomorphic = monomorphicLoop(iterations);
                System.out.println("phase=bimorphic");
                long bimorphic = bimorphicLoop(iterations);
                System.out.println("mode=lifecycle checksum=" + (monomorphic + bimorphic));
            }
            case "escape" -> {
                int iterations = parsePositive(args, 1, 50_000_000);
                System.out.println("mode=escape checksum=" + allocationLoop(iterations));
            }
            case "hold" -> {
                int iterations = parsePositive(args, 1, 5_000_000);
                if (args.length <= 2) {
                    throw new IllegalArgumentException("hold mode requires a release-file path");
                }
                Path releaseFile = Path.of(args[2]).toAbsolutePath();
                int timeoutMillis = parsePositive(args, 3, 60_000);
                long checksum = semanticLoop(iterations);
                System.out.println("mode=hold checksum=" + checksum);
                System.out.println("pid=" + ProcessHandle.current().pid() + " ready=codecache");
                long deadline = System.nanoTime() + timeoutMillis * 1_000_000L;
                while (!Files.exists(releaseFile)) {
                    if (System.nanoTime() >= deadline) {
                        throw new IllegalStateException("release file did not arrive before timeout");
                    }
                    Thread.sleep(50);
                }
                System.out.println("release=observed");
            }
            case "codecache" -> {
                int classCount = parsePositive(args, 1, 4_000);
                int callsPerClass = parsePositive(args, 2, 1_000);
                long checksum = codeCachePressure(classCount, callsPerClass);
                System.out.println("mode=codecache classes=" + classCount + " checksum=" + checksum);
            }
            default -> throw new IllegalArgumentException("unknown mode: " + mode);
        }
    }
}