package dst.ass2.di.agent;

import dst.ass2.di.annotation.Component;
import javassist.*;

import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.IllegalClassFormatException;
import java.lang.instrument.Instrumentation;
import java.security.ProtectionDomain;

public class InjectorAgent implements ClassFileTransformer {
    public static void premain(String args, Instrumentation instrumentation) {
        instrumentation.addTransformer(new InjectorAgent());
    }

    @Override
    public byte[] transform(ClassLoader loader, String className,
                            Class<?> classBeingRedefined, ProtectionDomain protectionDomain,
                            byte[] classfileBuffer) throws IllegalClassFormatException {

        ClassPool classPool = new ClassPool();
        classPool.appendClassPath(new LoaderClassPath(loader));
        try {
            CtClass ctClass = classPool.makeClass(new ByteArrayInputStream(classfileBuffer));


            // only @Components may continue.
            if (!ctClass.hasAnnotation(Component.class)) {
                return classfileBuffer;
            }

            // insert our code into constructors
            for (CtConstructor ctor : ctClass.getConstructors()) {
                try {
                    ctor.insertAfter("dst.ass2.di.InjectionControllerFactory.getTransparentInstance().initialize(this);");
                } catch (CannotCompileException e) {
                    e.printStackTrace();
                    return classfileBuffer;
                }
            }

            // output new bytecode.
            return ctClass.toBytecode();
        } catch (IOException | CannotCompileException e) {
            e.printStackTrace();
            return classfileBuffer;
        }
    }
}
