package com.hypnoticocelot.telemetry.agent;

import com.hypnoticocelot.telemetry.agent.handlers.MethodInstrumentationHandler;
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.HashSet;
import java.util.Iterator;
import java.util.Set;

public class TelemetryTransformer implements ClassFileTransformer {
    private final Set<MethodInstrumentationHandler> handlers = new HashSet<>();

    public void addHandler(MethodInstrumentationHandler handler) {
        handlers.add(handler);
    }

    @Override
    public byte[] transform(ClassLoader loader,
                            String className,
                            Class<?> classBeingRedefined,
                            ProtectionDomain protectionDomain,
                            byte[] classfileBuffer) throws IllegalClassFormatException {
        try {
            ClassReader cr = new ClassReader(classfileBuffer);
            ClassNode classNode = new ClassNode();
            cr.accept(classNode, ClassReader.EXPAND_FRAMES);

            boolean classUpdated = false;
            final Iterable<MethodNode> methodNodes = classNode.methods;
            for (MethodNode methodNode : methodNodes) {
                for (MethodInstrumentationHandler handler : handlers) {
                    if (handler.likes(classNode, methodNode)) {
                        System.out.println("Instrumenting method: handler=" + handler.getClass().getName() +
                                "; class=" + classNode.name +
                                "; method=" + methodNode.name);
                        InsnList beginList = new InsnList();
                        beginList.add(new FieldInsnNode(Opcodes.GETSTATIC, "java/lang/System", "out", "Ljava/io/PrintStream;"));
                        beginList.add(new LdcInsnNode("[BEGIN] YOUR ASS HAS BEEN INSTRUMENTED! " + methodNode.name));
                        beginList.add(new MethodInsnNode(Opcodes.INVOKEVIRTUAL, "java/io/PrintStream", "println", "(Ljava/lang/String;)V"));

                        methodNode.instructions.insert(beginList);

                        Iterator<AbstractInsnNode> insnNodes = methodNode.instructions.iterator();
                        while (insnNodes.hasNext()) {
                            AbstractInsnNode insn = insnNodes.next();

                            if (insn.getOpcode() == Opcodes.IRETURN
                                    || insn.getOpcode() == Opcodes.RETURN
                                    || insn.getOpcode() == Opcodes.ARETURN
                                    || insn.getOpcode() == Opcodes.LRETURN
                                    || insn.getOpcode() == Opcodes.DRETURN) {
                                InsnList endList = new InsnList();
                                endList.add(new FieldInsnNode(Opcodes.GETSTATIC, "java/lang/System", "out", "Ljava/io/PrintStream;"));
                                endList.add(new LdcInsnNode("[END] YOUR ASS HAS BEEN INSTRUMENTED! " + methodNode.name));
                                endList.add(new MethodInsnNode(Opcodes.INVOKEVIRTUAL, "java/io/PrintStream", "println", "(Ljava/lang/String;)V"));
                                methodNode.instructions.insertBefore(insn, endList);
                            }
                        }

                        classUpdated = true;
                        break;
                    }
                }
            }

            if (classUpdated) {
                ClassWriter cw = new ClassWriter(ClassWriter.COMPUTE_MAXS | ClassWriter.COMPUTE_FRAMES);
                classNode.accept(cw);

                return cw.toByteArray();
            } else {
                return null;
            }
        } catch (Throwable t) {
            t.printStackTrace();
            throw t;
        }
    }
}

@@ -148,91 +148,9 @@
 asm.
-ClassReader;%0Aimport org.objectweb.asm.ClassWriter;%0Aimport org.objectweb.asm.Opcodes
+*
 ;%0Aim
@@ -1801,37 +1801,35 @@
 ginList.add(new 
-Field
+Ldc
 InsnNode(Opcodes
@@ -1825,76 +1825,14 @@
 ode(
-Opcodes.GETSTATIC, %22java/lang/System%22, %22out%22, %22Ljava/io/PrintStream;
+%22SPAN!
 %22));
@@ -1866,35 +1866,38 @@
 ginList.add(new 
-Ldc
+Method
 InsnNode(%22%5BBEGIN
@@ -1893,68 +1893,149 @@
 ode(
-%22%5BBEGIN%5D YOUR ASS HAS BEEN INSTRUMENTED! %22 + methodNode.name
+Opcodes.INVOKESTATIC, %22com/hypnoticocelot/telemetry/tracing/Span%22, %22start%22, %22(Ljava/lang/String;)Lcom/hypnoticocelot/telemetry/tracing/Span;%22
 ));%0A
@@ -2068,38 +2068,35 @@
 ginList.add(new 
-Method
+Var
 InsnNode(Opcodes
@@ -2100,80 +2100,40 @@
 des.
-INVOKEVIRTUAL, %22java/io/PrintStream%22, %22println%22, %22(Ljava/lang/String;)V%22
+ASTORE, methodNode.maxLocals + 1
 ));%0A
@@ -2914,13 +2914,11 @@
 new 
-Field
+Var
 Insn
@@ -2934,191 +2934,39 @@
 des.
-GETSTATIC, %22java/lang/System%22, %22out%22, %22Ljava/io/PrintStream;%22));%0A                                endList.add(new LdcInsnNode(%22%5BEND%5D YOUR ASS HAS BEEN INSTRUMENTED! %22 + methodNode.name
+ALOAD, methodNode.maxLocals + 1
 ));%0A
@@ -3056,61 +3056,63 @@
 L, %22
-java/io/PrintStream%22, %22println%22, %22(Ljava/lang/String;
+com/hypnoticocelot/telemetry/tracing/Span%22, %22close%22, %22(
 )V%22)
@@ -3654,16 +3654,16 @@
 n null;%0A
-
         
@@ -3660,32 +3660,122 @@
 ;%0A            %7D%0A
+        %7D catch (Exception e) %7B%0A            e.printStackTrace();%0A            return null;%0A
         %7D catch 

Work in progress commit of the agent. This may currently be broken due to the addition of SpanData.