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.security.ProtectionDomain;
import java.util.Arrays;
import java.util.List;

public class ClassTransformer implements ClassFileTransformer {

    private static final List<String> classBlacklist = Arrays.asList(
            "java/",
            "javax/",
            "sun/",
            "com/sun/",
            "jdk/",
            "ch/usi/inf/sp/dbi/profiler/"
    );

    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) {

        if (classBlacklist.stream().anyMatch(className::startsWith)) {
            System.out.printf("Skipping class <%s, %s> (in blacklist)%n", loader, className);
            return classfileBuffer;
        }

        System.out.printf("About to transform class <%s, %s>%n", loader, className);

        try {
            return instrument(classfileBuffer);
        } catch (Throwable e) {
            System.out.printf("Error instrumenting class <%s>: %s%n", className, e.getMessage());
            e.printStackTrace(System.err);

            return classfileBuffer; // return non-instrumented class
        }
    }
}
