Custom screen builder for main config

lib/src/main/java/me/lortseam/completeconfig/gui/ConfigScreenBuilder.java

@@ -5,8 +5,22 @@ import net.fabricmc.api.EnvType;
 import net.fabricmc.api.Environment;
 import net.minecraft.client.gui.screen.Screen;
 
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Optional;
+
 public abstract class ConfigScreenBuilder {
 
+    private static final Map<String, ConfigScreenBuilder> mainBuilders = new HashMap<>();
+
+    public static void setMainBuilder(String modID, ConfigScreenBuilder screenBuilder) {
+        mainBuilders.put(modID, screenBuilder);
+    }
+
+    public static Optional<ConfigScreenBuilder> getMainBuilder(String modID) {
+        return Optional.ofNullable(mainBuilders.get(modID));
+    }
+
     /**
      * Builds a screen based on a config.
      *

lib/src/main/java/me/lortseam/completeconfig/gui/ModMenuIntegration.java

@@ -10,11 +10,13 @@ import java.util.Map;
 
 public final class ModMenuIntegration implements ModMenuApi {
 
-    private final ConfigScreenBuilder screenBuilder = new ClothConfigScreenBuilder();
+    private final ConfigScreenBuilder defaultScreenBuilder = new ClothConfigScreenBuilder();
 
     @Override
     public Map<String, ConfigScreenFactory<?>> getProvidedConfigScreenFactories() {
-        return Maps.transformValues(Config.getMainConfigs(), config -> parentScreen -> screenBuilder.build(parentScreen, config));
+        return Maps.transformValues(Config.getMainConfigs(), config -> parentScreen -> {
+            return ConfigScreenBuilder.getMainBuilder(config.getModID()).orElse(defaultScreenBuilder).build(parentScreen, config);
+        });
     }
 
 }