/*

 * Copyright (c) 2014, Oracle and/or its affiliates. All rights reserved.

 * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.

 *

 * This code is free software; you can redistribute it and/or modify it

 * under the terms of the GNU General Public License version 2 only, as

 * published by the Free Software Foundation.

 *

 * This code is distributed in the hope that it will be useful, but WITHOUT

 * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or

 * FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License

 * version 2 for more details (a copy is included in the LICENSE file that

 * accompanied this code).

 *

 * You should have received a copy of the GNU General Public License version

 * 2 along with this work; if not, write to the Free Software Foundation,

 * Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA.

 *

 * Please contact Oracle, 500 Oracle Parkway, Redwood Shores, CA 94065 USA

 * or visit www.oracle.com if you need additional information or have any

 * questions.

 */

package asmlib;

import java.io.PrintStream;

import java.util.HashSet;

import java.util.Set;

import java.util.concurrent.atomic.AtomicInteger;

import jdk.internal.org.objectweb.asm.ClassReader;

import jdk.internal.org.objectweb.asm.ClassVisitor;

import jdk.internal.org.objectweb.asm.ClassWriter;

import jdk.internal.org.objectweb.asm.MethodVisitor;

import jdk.internal.org.objectweb.asm.Opcodes;

import java.util.function.Consumer;

import jdk.internal.org.objectweb.asm.Type;

public class Instrumentor {

    public static class InstrHelper {

        private final MethodVisitor mv;

        private final String name;

        InstrHelper(MethodVisitor mv, String name) {

            this.mv = mv;

            this.name = name;

        }

        public String getName() {

            return this.name;

        }

        public void invokeStatic(String owner, String name, String desc, boolean itf) {

            mv.visitMethodInsn(Opcodes.INVOKESTATIC, owner, name, desc, itf);

        }

        public void invokeSpecial(String owner, String name, String desc) {

            mv.visitMethodInsn(Opcodes.INVOKESPECIAL, owner, name, desc, false);

        }

        public void invokeVirtual(String owner, String name, String desc) {

            mv.visitMethodInsn(Opcodes.INVOKEVIRTUAL, owner, name, desc, false);

        }

        public void push(int val) {

            if (val >= -1 && val <= 5) {

                mv.visitInsn(Opcodes.ICONST_0 + val);

            } else if (val >= Byte.MIN_VALUE && val <= Byte.MAX_VALUE) {

                mv.visitIntInsn(Opcodes.BIPUSH, val);

            } else if (val >= Short.MIN_VALUE && val <= Short.MAX_VALUE) {

                mv.visitIntInsn(Opcodes.SIPUSH, val);

            } else {

                mv.visitLdcInsn(val);

            }

        }

        public void push(Object val) {

            mv.visitLdcInsn(val);

        }

        public void println(String s) {

            mv.visitFieldInsn(Opcodes.GETSTATIC, Type.getInternalName(System.class), "out", Type.getDescriptor(PrintStream.class));

            mv.visitLdcInsn(s);

            mv.visitMethodInsn(Opcodes.INVOKEVIRTUAL, Type.getInternalName(PrintStream.class), "println", Type.getMethodDescriptor(Type.VOID_TYPE, Type.getType(String.class)), false);

        }

    }

    public static Instrumentor instrFor(byte[] classData) {

        return new Instrumentor(classData);

    }

    private final ClassReader cr;

    private final ClassWriter output;

    private ClassVisitor instrumentingVisitor = null;

    private final AtomicInteger matches = new AtomicInteger(0);

    private Instrumentor(byte[] classData) {

        cr = new ClassReader(classData);

        output = new ClassWriter(ClassWriter.COMPUTE_MAXS);

        instrumentingVisitor = output;

    }

    public synchronized Instrumentor addMethodEntryInjection(String methodName, Consumer<InstrHelper> injector) {

        instrumentingVisitor = new ClassVisitor(Opcodes.ASM5, instrumentingVisitor) {

            @Override

            public MethodVisitor visitMethod(int access, String name, String desc, String signature, String[] exceptions) {

                MethodVisitor mv = super.visitMethod(access, name, desc, signature, exceptions);

                if (name.equals(methodName)) {

                    matches.getAndIncrement();

                    mv = new MethodVisitor(Opcodes.ASM5, mv) {

                        @Override

                        public void visitCode() {

                            injector.accept(new InstrHelper(mv, name));

                        }

                    };

                }

                return mv;

            }

        };

        return this;

    }

    public synchronized Instrumentor addNativeMethodTrackingInjection(String prefix, Consumer<InstrHelper> injector) {

        instrumentingVisitor = new ClassVisitor(Opcodes.ASM5, instrumentingVisitor) {

            private final Set<Consumer<ClassVisitor>> wmGenerators = new HashSet<>();

            private String className;

            @Override

            public void visit(int version, int access, String name, String signature, String superName, String[] interfaces) {

                this.className = name;

                super.visit(version, access, name, signature, superName, interfaces);

            }

            @Override

            public MethodVisitor visitMethod(int access, String name, String desc, String signature, String[] exceptions) {

                if ((access & Opcodes.ACC_NATIVE) != 0) {

                    matches.getAndIncrement();

                    String newName = prefix + name;

                    wmGenerators.add((v)->{

                        MethodVisitor mv = v.visitMethod(access & ~Opcodes.ACC_NATIVE, name, desc, signature, exceptions);

                        mv.visitCode();

                        injector.accept(new InstrHelper(mv, name));

                        Type[] argTypes = Type.getArgumentTypes(desc);

                        Type retType = Type.getReturnType(desc);

                        boolean isStatic = (access & Opcodes.ACC_STATIC) != 0;

                        if (!isStatic) {

                            mv.visitIntInsn(Opcodes.ALOAD, 0); // load "this"

                        }

                        // load the method parameters

                        if (argTypes.length > 0) {

                            int ptr = isStatic ? 0 : 1;

                            for(Type argType : argTypes) {

                                mv.visitIntInsn(argType.getOpcode(Opcodes.ILOAD), ptr);

                                ptr += argType.getSize();

                            }

                        }

                        mv.visitMethodInsn(isStatic ? Opcodes.INVOKESTATIC : Opcodes.INVOKESPECIAL, className, newName, desc, false);

                        mv.visitInsn(retType.getOpcode(Opcodes.IRETURN));

                        mv.visitMaxs(1, 1); // dummy call; let ClassWriter to deal with this

                        mv.visitEnd();

                    });

                    return super.visitMethod(access, newName, desc, signature, exceptions);

                }

                return super.visitMethod(access, name, desc, signature, exceptions);

            }

            @Override

            public void visitEnd() {

                wmGenerators.stream().forEach((e) -> {

                    e.accept(cv);

                });

                super.visitEnd();

            }

        };

        return this;

    }

    public synchronized byte[] apply() {

        cr.accept(instrumentingVisitor, ClassReader.SKIP_DEBUG + ClassReader.EXPAND_FRAMES);

        return matches.get() == 0 ? null : output.toByteArray();

    }

}