import java.io.InputStream;
import java.lang.classfile.ClassFile;
import java.lang.classfile.CodeModel;
import java.lang.classfile.MethodModel;
import java.lang.classfile.instruction.InvokeInstruction;
import java.lang.classfile.instruction.ConstantInstruction;
import java.lang.invoke.MethodType;
import java.lang.reflect.Method;
import java.util.*;
import java.util.jar.JarEntry;
import java.util.jar.JarFile;

/**
 * Static scan of jars for call sites of the sun.misc.Unsafe methods that are deprecated for removal (JEP 471).
 *
 * It reads bytecode with the java.lang.classfile API (final in JDK 24). It does NOT run anything, so it reports
 * "this jar contains a call to X", never "this call is reached at run time". Which methods count as
 * "memory-access" is decided by asking the running JDK: a method is counted when its own Unsafe overload carries
 * @Deprecated(forRemoval = true).
 *
 * usage: java ScanJar.java [--classes] jar [jar ...]
 */
public class ScanJar {
    static final String UNSAFE = "sun/misc/Unsafe";

    public static void main(String[] args) throws Exception {
        boolean perClass = false;
        List<String> jars = new ArrayList<>();
        for (String a : args) { if (a.equals("--classes")) perClass = true; else jars.add(a); }

        Set<String> removal = new HashSet<>();       // name + descriptor of every forRemoval overload
        int total = 0;
        for (Method m : Class.forName("sun.misc.Unsafe").getDeclaredMethods()) {
            if (!java.lang.reflect.Modifier.isPublic(m.getModifiers())) continue;
            total++;
            Deprecated d = m.getAnnotation(Deprecated.class);
            if (d != null && d.forRemoval()) {
                removal.add(m.getName() + MethodType.methodType(m.getReturnType(), m.getParameterTypes()).descriptorString());
            }
        }
        System.out.println("this JDK: sun.misc.Unsafe has " + total + " public methods, " + removal.size() + " deprecated for removal");

        for (String jar : jars) scan(jar, removal, perClass);
    }

    static void scan(String jar, Set<String> removal, boolean perClass) throws Exception {
        int classes = 0, refClasses = 0, callSites = 0, mrClasses = 0, stringRefs = 0;
        Map<String, Integer> methods = new TreeMap<>();
        Map<String, Integer> perClassCalls = new TreeMap<>();
        Set<String> otherUnsafeCalls = new TreeSet<>();
        Set<String> mrCallers = new TreeSet<>();
        try (JarFile jf = new JarFile(jar)) {
            for (JarEntry e : Collections.list(jf.entries())) {
                String n = e.getName();
                if (!n.endsWith(".class") || n.endsWith("module-info.class")) continue;
                boolean mr = n.startsWith("META-INF/versions/");
                byte[] bytes;
                try (InputStream in = jf.getInputStream(e)) { bytes = in.readAllBytes(); }
                var cm = ClassFile.of().parse(bytes);
                classes++;
                boolean ref = false; int calls = 0;
                for (MethodModel mm : cm.methods()) {
                    CodeModel code = mm.code().orElse(null);
                    if (code == null) continue;
                    for (var el : code) {
                        if (el instanceof InvokeInstruction ii && ii.owner().asInternalName().equals(UNSAFE)) {
                            ref = true;
                            String key = ii.name().stringValue() + ii.typeSymbol().descriptorString();
                            if (removal.contains(key)) { calls++; methods.merge(ii.name().stringValue(), 1, Integer::sum); }
                            else otherUnsafeCalls.add(ii.name().stringValue());
                        } else if (el instanceof ConstantInstruction.LoadConstantInstruction lc
                                   && lc.constantEntry().constantValue() instanceof String s
                                   && (s.equals("sun.misc.Unsafe") || s.equals("theUnsafe"))) {
                            stringRefs++; ref = true;
                        }
                    }
                }
                if (ref) refClasses++;
                if (calls > 0) {
                    callSites += calls;
                    String cn = n.substring(0, n.length() - 6).replace('/', '.');
                    perClassCalls.put(cn, calls);
                    if (mr) { mrClasses++; mrCallers.add(cn); }
                }
            }
        }
        String name = jar.substring(jar.lastIndexOf('/') + 1);
        System.out.println();
        System.out.println(name + ": " + classes + " classes scanned");
        System.out.println("  classes touching sun.misc.Unsafe (call, or the strings 'sun.misc.Unsafe'/'theUnsafe'): " + refClasses);
        System.out.println("  classes calling a removal-deprecated Unsafe method: " + perClassCalls.size() + "  (" + callSites + " call sites, " + methods.size() + " distinct methods)");
        if (!methods.isEmpty()) System.out.println("  methods: " + methods);
        if (!otherUnsafeCalls.isEmpty()) System.out.println("  other (not deprecated for removal) Unsafe methods called: " + otherUnsafeCalls);
        if (!mrCallers.isEmpty()) System.out.println("  of those classes, under META-INF/versions: " + mrCallers.size());
        if (perClass) perClassCalls.forEach((k, v) -> System.out.println("    " + k + " " + v));
    }
}