Minestom/code-generators/src/main/java/net/minestom/codegen/entity/VillagerProfessionGenerator.java

package net.minestom.codegen.entity;

import com.google.gson.JsonArray;
import com.google.gson.JsonElement;
import com.google.gson.JsonObject;
import com.squareup.javapoet.*;
import net.minestom.codegen.MinestomCodeGenerator;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.lang.model.element.Modifier;
import java.io.File;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.util.List;

public final class VillagerProfessionGenerator extends MinestomCodeGenerator {
    private static final Logger LOGGER = LoggerFactory.getLogger(VillagerProfessionGenerator.class);
    private final InputStream villagerProfessionsFile;
    private final File outputFolder;


    public VillagerProfessionGenerator(@Nullable InputStream villagerProfessionsFile, @NotNull File outputFolder) {
        this.villagerProfessionsFile = villagerProfessionsFile;
        this.outputFolder = outputFolder;
    }

    @Override
    public void generate() {
        if (villagerProfessionsFile == null) {
            LOGGER.error("Failed to find villager_professions.json.");
            LOGGER.error("Stopped code generation for villager professions.");
            return;
        }
        if (!outputFolder.exists() && !outputFolder.mkdirs()) {
            LOGGER.error("Output folder for code generation does not exist and could not be created.");
            return;
        }
        // Important classes we use alot
        ClassName namespaceIDClassName = ClassName.get("net.minestom.server.utils", "NamespaceID");
        ClassName rawVillagerProfessionDataClassName = ClassName.get("net.minestom.server.raw_data", "RawVillagerProfessionData");
        ClassName registryClassName = ClassName.get("net.minestom.server.registry", "Registry");

        JsonArray villagerProfessions = GSON.fromJson(new InputStreamReader(villagerProfessionsFile), JsonArray.class);
        ClassName villagerProfessionClassName = ClassName.get("net.minestom.server.entity.metadata.villager", "VillagerProfession");

        // Particle
        TypeSpec.Builder villagerProfessionClass = TypeSpec.classBuilder(villagerProfessionClassName)
                .addSuperinterface(ClassName.get("net.kyori.adventure.key", "Keyed"))
                .addModifiers(Modifier.PUBLIC).addJavadoc("AUTOGENERATED by " + getClass().getSimpleName());
        villagerProfessionClass.addField(
                FieldSpec.builder(namespaceIDClassName, "id")
                        .addModifiers(Modifier.PRIVATE, Modifier.FINAL).addAnnotation(NotNull.class).build()
        );
        villagerProfessionClass.addField(
                FieldSpec.builder(rawVillagerProfessionDataClassName, "villagerProfessionData")
                        .addModifiers(Modifier.PRIVATE, Modifier.VOLATILE)
                        .addAnnotation(NotNull.class)
                        .build()
        );
        villagerProfessionClass.addMethod(
                MethodSpec.constructorBuilder()
                        .addParameter(ParameterSpec.builder(namespaceIDClassName, "id").addAnnotation(NotNull.class).build())
                        .addParameter(ParameterSpec.builder(rawVillagerProfessionDataClassName, "villagerProfessionData").addAnnotation(NotNull.class).build())
                        .addStatement("this.id = id")
                        .addStatement("this.villagerProfessionData = villagerProfessionData")
                        .addModifiers(Modifier.PROTECTED)
                        .build()
        );
        // Override key method (adventure)
        villagerProfessionClass.addMethod(
                MethodSpec.methodBuilder("key")
                        .returns(ClassName.get("net.kyori.adventure.key", "Key"))
                        .addAnnotation(Override.class)
                        .addAnnotation(NotNull.class)
                        .addStatement("return this.id")
                        .addModifiers(Modifier.PUBLIC)
                        .build()
        );
        // getId method
        villagerProfessionClass.addMethod(
                MethodSpec.methodBuilder("getId")
                        .returns(namespaceIDClassName)
                        .addAnnotation(NotNull.class)
                        .addStatement("return this.id")
                        .addModifiers(Modifier.PUBLIC)
                        .build()
        );
        // getVillagerProfessionData method
        villagerProfessionClass.addMethod(
                MethodSpec.methodBuilder("getVillagerProfessionData")
                        .returns(rawVillagerProfessionDataClassName)
                        .addAnnotation(NotNull.class)
                        .addStatement("return this.villagerProfessionData")
                        .addModifiers(Modifier.PUBLIC, Modifier.FINAL)
                        .build()
        );
        // setVillagerProfessionData method
        villagerProfessionClass.addMethod(
                MethodSpec.methodBuilder("setVillagerProfessionData")
                        .addParameter(ParameterSpec.builder(rawVillagerProfessionDataClassName, "villagerProfessionData").addAnnotation(NotNull.class).build())
                        .addStatement("this.villagerProfessionData = villagerProfessionData")
                        .addModifiers(Modifier.PUBLIC, Modifier.FINAL)
                        .build()
        );
        // getNumericalId
        villagerProfessionClass.addMethod(
                MethodSpec.methodBuilder("getNumericalId")
                        .returns(TypeName.INT)
                        .addStatement("return $T.VILLAGER_PROFESSION_REGISTRY.getId(this)", registryClassName)
                        .addModifiers(Modifier.PUBLIC)
                        .build()
        );
        // fromId Method
        villagerProfessionClass.addMethod(
                MethodSpec.methodBuilder("fromId")
                        .returns(villagerProfessionClassName)
                        .addAnnotation(Nullable.class)
                        .addParameter(TypeName.INT, "id")
                        .addStatement("return $T.VILLAGER_PROFESSION_REGISTRY.get((short) id)", registryClassName)
                        .addModifiers(Modifier.PUBLIC, Modifier.STATIC)
                        .build()
        );
        // fromId Method
        villagerProfessionClass.addMethod(
                MethodSpec.methodBuilder("fromId")
                        .returns(villagerProfessionClassName)
                        .addAnnotation(NotNull.class)
                        .addParameter(ClassName.get("net.kyori.adventure.key", "Key"), "id")
                        .addStatement("return $T.VILLAGER_PROFESSION_REGISTRY.get(id)", registryClassName)
                        .addModifiers(Modifier.PUBLIC, Modifier.STATIC)
                        .build()
        );
        // toString method
        villagerProfessionClass.addMethod(
                MethodSpec.methodBuilder("toString")
                        .addAnnotation(NotNull.class)
                        .addAnnotation(Override.class)
                        .returns(String.class)
                        // this resolves to [Namespace]
                        .addStatement("return \"[\" + this.id + \"]\"")
                        .addModifiers(Modifier.PUBLIC)
                        .build()
        );
        // values method
        villagerProfessionClass.addMethod(
                MethodSpec.methodBuilder("values")
                        .addAnnotation(NotNull.class)
                        .returns(ParameterizedTypeName.get(ClassName.get(List.class), villagerProfessionClassName))
                        .addStatement("return $T.VILLAGER_PROFESSION_REGISTRY.values()", registryClassName)
                        .addModifiers(Modifier.PUBLIC, Modifier.STATIC)
                        .build()
        );
        CodeBlock.Builder staticBlock = CodeBlock.builder();
        // Use data
        for (JsonElement vp : villagerProfessions) {
            JsonObject villagerProfession = vp.getAsJsonObject();

            String villagerProfessionName = villagerProfession.get("name").getAsString();
            JsonElement workSound = villagerProfession.get("workSound");
            if (workSound == null) {
                villagerProfessionClass.addField(
                        FieldSpec.builder(
                                villagerProfessionClassName,
                                villagerProfessionName
                        ).initializer(
                                "new $T($T.from($S), new $T(() -> null))",
                                villagerProfessionClassName,
                                namespaceIDClassName,
                                villagerProfession.get("id").getAsString(),

                                rawVillagerProfessionDataClassName
                        ).addModifiers(Modifier.PUBLIC, Modifier.STATIC, Modifier.FINAL).build()
                );
            } else {
                villagerProfessionClass.addField(
                        FieldSpec.builder(
                                villagerProfessionClassName,
                                villagerProfessionName
                        ).initializer(
                                "new $T($T.from($S), new $T(() -> $T.SOUND_EVENT_REGISTRY.get($S)))",
                                villagerProfessionClassName,
                                namespaceIDClassName,
                                villagerProfession.get("id").getAsString(),

                                rawVillagerProfessionDataClassName,
                                registryClassName,
                                workSound.getAsString()
                        ).addModifiers(Modifier.PUBLIC, Modifier.STATIC, Modifier.FINAL).build()
                );
            }

            // Add to static init.
            staticBlock.addStatement("$T.VILLAGER_PROFESSION_REGISTRY.register($N)", registryClassName, villagerProfessionName);
        }

        villagerProfessionClass.addStaticBlock(staticBlock.build());

        // Write files to outputFolder
        writeFiles(
                List.of(
                        JavaFile.builder("net.minestom.server.entity.metadata.villager", villagerProfessionClass.build())
                                .indent("    ")
                                .skipJavaLangImports(true)
                                .build()
                ),
                outputFolder
        );
    }
}