Messing around with XJC plugins to get it to generate Protobuf Marshaller classes for Infinispan

ProtobufPlugin.java

package com.example;

import com.sun.codemodel.*;
import com.sun.tools.xjc.Options;
import com.sun.tools.xjc.Plugin;
import com.sun.tools.xjc.model.Model;
import com.sun.tools.xjc.outline.ClassOutline;
import com.sun.tools.xjc.outline.EnumConstantOutline;
import com.sun.tools.xjc.outline.EnumOutline;
import com.sun.tools.xjc.outline.Outline;
import org.infinispan.protostream.MessageMarshaller;
import org.infinispan.protostream.annotations.ProtoEnumValue;
import org.infinispan.protostream.annotations.ProtoField;
import org.xml.sax.ErrorHandler;

import java.io.IOException;
import java.util.*;
import java.util.logging.Logger;

/**
 * Created by tdonohue on 15/06/2018.
 */
public class ProtobufPlugin extends Plugin {

    private static final Logger LOG = Logger.getLogger(ProtobufPlugin.class.getName());

    public ProtobufPlugin() {
    }

    @Override
    public String getOptionName() {
        return "Xprotobuf";
    }

    @Override
    public String getUsage() {
        return "  -Xprotobuf           : Add protobuf annotations to all fields";
    }

    @Override
    public void postProcessModel(Model model, ErrorHandler errorHandler) {
        super.postProcessModel(model, errorHandler);
    }

    public boolean run(Outline outline, Options opt, ErrorHandler errorHandler) {
        LOG.info("Running Protobuf plugin.");

        JPackage targetPackage = null;
        ClassOutline targetClass = null;

        List<MarshallerOutline> marshallers = new ArrayList<MarshallerOutline>();

        /* Add @ProtoEnumValue annotation to enum values */
        for (EnumOutline enumOutline : outline.getEnums()) {
            int protoEnumCounter = 1;

            for (EnumConstantOutline eco : enumOutline.constants) {
                eco.constRef.annotate(ProtoEnumValue.class)
                        .param("number", protoEnumCounter);

                protoEnumCounter++;
            }

        }

        for (ClassOutline classOutline : outline.getClasses()) {

            int protoFieldCounter = 1;

            List<Map> fields = new ArrayList<Map>();

            Collection<JMethod> methods = classOutline.implClass.methods();

            for (JMethod method : methods) {
                if (method.name().startsWith("get")) {

                    Map<String, String> fieldParams = new HashMap<String, String>();

                    LOG.info("Adding annotation to " + method.name());

                    JAnnotationUse method2 = method.annotate(ProtoField.class)
                            .param("number", protoFieldCounter);

                    if (method.type().name().startsWith("List<")) {
                        LOG.info("This is a LIST class!");
                        JAnnotationUse method3 = method2.param("collectionImplementation",
                                ArrayList.class);
                    }

                    fieldParams.put("getter", method.name());

/*
                    List<Object> body = method.body().getContents();

//                    Object lastStatement = body.get(body.size() - 1);

                    LOG.info("Final statement in body is: " + body.get(body.size() - 1));

                    fields.add(fieldParams);
*/

                    protoFieldCounter++;
                }


            }
/*

            MarshallerOutline marshaller = new MarshallerOutline();
            marshaller.setTargetPackage(classOutline.implClass._package());
            marshaller.setClassOutline(classOutline);
            marshaller.setFields(fields);

            marshallers.add(marshaller);
*/


        }

/*
        for (MarshallerOutline marshaller : marshallers) {
            createClazz(outline, marshaller);
        }
*/

        return true;
    }

    private void createClazz(Outline outline, MarshallerOutline marshaller) {
        JCodeModel cm = outline.getCodeModel();

        JPackage targetPackage = marshaller.getTargetPackage();
        ClassOutline clazzOutline = marshaller.getClassOutline();

        JDefinedClass impl = outline.getClassFactory().createClass(
                targetPackage,
                clazzOutline.implClass.name() + "Marshaller",
                null);

        impl._implements(cm.ref(MessageMarshaller.class).narrow(clazzOutline.implClass));

        impl.method(JMod.PUBLIC, String.class, "getTypeName") // TODO add annotation
                .body()._return(JExpr.lit("book_sample." + clazzOutline.implClass.name())); // TODO - return correct package name

        impl.method(JMod.PUBLIC, cm.ref(Class.class).narrow(clazzOutline.implClass), "getJavaClass") // TODO - add annotation
                .body()._return(JExpr.dotclass(clazzOutline.implClass));

        JMethod writeTo = impl.method(JMod.PUBLIC, cm.VOID, "writeTo")
                ._throws(IOException.class);
        JVar writer = writeTo.param(MessageMarshaller.ProtoStreamWriter.class, "writer");
        JVar obj = writeTo.param(clazzOutline.implClass, "obj");
        JBlock writeToBody = writeTo.body();

        for (Map fieldParams : marshaller.getFields()) {
            LOG.info("Getter is: " + fieldParams.get("getter"));

            writeToBody.invoke(writer, "writeString")
                    .arg("FIELD_NAME")
                    .arg(obj.invoke(fieldParams.get("getter").toString()));
        }

//        writeToBody.invoke(writer, "writeString")
//                .arg("hiya")
//                .arg(obj.invoke("getTitle"));

                // body().invoke($var,"add").arg(createJAXBElement(loop.var()));

        JMethod method = impl.method(JMod.PUBLIC, clazzOutline.implClass, "readFrom");
        method.param(MessageMarshaller.ProtoStreamReader.class, "reader");
        method._throws(IOException.class);

//        JCodeModel codeModel = outline.getCodeModel();
//        codeModel.ref()
//
//        method.body().decl(String.class, "title");
//        method.body().decl(com.sun.codemodel.internal.JCodeModel.)

        LOG.info("Added a class: " + impl.toString());
    }

}