package com.github.oreissig.agentinjecteur.agent;

import com.github.oreissig.agentinjecteur.rt.MembersInjector;
import javassist.*;

import javax.inject.Inject;
import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.util.Arrays;
import java.util.List;

import static java.util.stream.Collectors.toList;
import static javassist.Modifier.FINAL;

class InjectionTransformer {

    private final ClassPool pool = ClassPool.getDefault();
    private final String membersInjector = MembersInjector.class.getName();

    public byte[] transform(String className, byte[] classfileBuffer) throws IOException, CannotCompileException {
        CtClass clazz = pool.makeClass(new ByteArrayInputStream(classfileBuffer));

        // TODO @Assisted for constructor params
        // TODO overwrite @FactoryMethod annotated static methods with getter
        List<CtField> injectTargets = Arrays.stream(clazz.getDeclaredFields())
                .filter(f -> f.hasAnnotation(Inject.class)).collect(toList());
        if (injectTargets.isEmpty()) {
            // do nothing
            return classfileBuffer;
        } else {
            injectMembers(clazz, injectTargets);
            return clazz.toBytecode();
        }
    }

    private void injectMembers(CtClass clazz, List<CtField> injectTargets) {
        try {
            for (CtField f : injectTargets) {
                // make field final
                // TODO check that there is no write to an injected field
                f.setModifiers(f.getModifiers() | FINAL);
                // TODO replace @javax.inject.Inject by proprietary @Injected annotation to avoid double injection

                String typeName = f.getType().getName();
                String init = "(" + typeName + ") " + membersInjector + ".get(" + typeName + ".class)";
                clazz.removeField(f);
                clazz.addField(f, init);
            }
        } catch (NotFoundException | CannotCompileException e) {
            throw new RuntimeException(e);
        }
    }
}

@@ -125,16 +125,85 @@
 ssist.*;
+%0Aimport javassist.expr.ExprEditor;%0Aimport javassist.expr.FieldAccess;
 %0A%0Aimport
@@ -427,16 +427,58 @@
 r.FINAL;
+%0Aimport static javassist.Modifier.PRIVATE;
 %0A%0Aclass 
@@ -1500,158 +1500,20 @@
     
- //
  make
- field final%0A                // TODO check that there is no write to an injected field%0A                f.setModifiers(f.getModifiers() %7C FINAL
+Final(f
 );%0A 
@@ -1942,24 +1942,24 @@
 eption e) %7B%0A
-
             
@@ -2004,11 +2004,900 @@
  %7D%0A    %7D
+%0A%0A    private void makeFinal(CtField f) throws CannotCompileException %7B%0A        if ((f.getModifiers() & PRIVATE) %3E 0) %7B%0A            // check that there is no write to an injected field%0A            f.getDeclaringClass().instrument(new ExprEditor() %7B%0A                public void edit(FieldAccess fa) throws CannotCompileException %7B%0A                    try %7B%0A                        if (f.equals(fa.getField()) && fa.isWriter()) %7B%0A                            throw new CannotCompileException(%22write to private injected field %22 +%0A                                    f.getName() + %22 not allowed%22);%0A                        %7D%0A                    %7D catch (NotFoundException e) %7B%0A                        throw new CannotCompileException(e);%0A                    %7D%0A                %7D%0A            %7D);%0A            // make field final%0A            f.setModifiers(f.getModifiers() %7C FINAL);%0A        %7D%0A    %7D
 %0A%7D%0A

check that there is no write to an injected field