package ch.usi.inf.sp.dbi.agent;

import org.objectweb.asm.ClassReader;
import org.objectweb.asm.ClassWriter;
import org.objectweb.asm.Opcodes;
import org.objectweb.asm.tree.*;

import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.IllegalClassFormatException;
import java.security.ProtectionDomain;
import java.util.Set;

public class ClassTransformer implements ClassFileTransformer {
    public static byte[] instrument(final byte[] bytecode) {
        final ClassReader cr = new ClassReader(bytecode);
        final ClassNode cn = new ClassNode();
        cr.accept(cn, ClassReader.SKIP_FRAMES);

        instrumentClassNode(cn);

        final ClassWriter cw = new ClassWriter(ClassWriter.COMPUTE_FRAMES);
        cn.accept(cw);
        return cw.toByteArray();
    }

    private static void instrumentClassNode(final ClassNode cn) {
        for (final MethodNode mn : cn.methods) {
            instrumentEntry(cn, mn);
            instrumentExit(cn, mn);
        }

    }

    private static void instrumentEntry(ClassNode cn, MethodNode mn) {
        InsnList patch = new InsnList(); // instructions to be added
        patch.add(new LdcInsnNode(cn.name));
        patch.add(new LdcInsnNode(mn.name));
        patch.add(new LdcInsnNode(mn.desc));
        patch.add(new MethodInsnNode(Opcodes.INVOKESTATIC,
                "ch/usi/inf/sp/dbi/profiler/Profiler",
                "entry",
                "(Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;)V"));

        mn.instructions.insert(patch);
    }

    private static void instrumentExit(ClassNode cn, MethodNode mn) {
        for (AbstractInsnNode i : mn.instructions) {

            if (i.getOpcode() >= Opcodes.IRETURN && i.getOpcode() <= Opcodes.RETURN) {
                InsnList patch = new InsnList(); // instructions to be added
                patch.add(new LdcInsnNode(cn.name));
                patch.add(new LdcInsnNode(mn.name));
                patch.add(new LdcInsnNode(mn.desc));
                patch.add(new MethodInsnNode(Opcodes.INVOKESTATIC,
                        "ch/usi/inf/sp/dbi/profiler/Profiler",
                        "exit",
                        "(Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;)V"));

                mn.instructions.insertBefore(i, patch);
            }
        }
    }

    @Override
    public byte[] transform(ClassLoader loader,
                            String className,
                            Class<?> classBeingRedefined,
                            ProtectionDomain protectionDomain,
                            byte[] classfileBuffer) throws IllegalClassFormatException {
//        className.startsWith("java/") ||
//                className.startsWith("javax/") ||
//                className.startsWith("sun/") ||
//                className.startsWith("com/sun/") ||
//                className.startsWith("jdk/") ||

        if (!className.startsWith("ch/usi/inf/sp/dbi/") ||
                className.startsWith("ch/usi/inf/sp/dbi/agent/") ||
                className.startsWith("ch/usi/inf/sp/dbi/profiler/")) {
            System.out.println("Skipping class <" + loader + ", " + className + ">");
            return classfileBuffer;
        }

        System.out.println("About to transform class <" + loader + ", " + className + ">");
        try {
            return instrument(classfileBuffer);
        } catch (Throwable e) {
            e.printStackTrace(System.err);
            return classfileBuffer;
        }
    }
}
