Tidy up mod lifecycle

1 files changed, 30 insertions, 2 deletions

diff --git a/spark-fabric/src/main/java/me/lucko/spark/fabric/FabricSparkMod.java b/spark-fabric/src/main/java/me/lucko/spark/fabric/FabricSparkMod.java
index 1dd4fd8..fdb359e 100644
--- a/spark-fabric/src/main/java/me/lucko/spark/fabric/FabricSparkMod.java
+++ b/spark-fabric/src/main/java/me/lucko/spark/fabric/FabricSparkMod.java
@@ -20,14 +20,19 @@
 
 package me.lucko.spark.fabric;
 
+import com.mojang.brigadier.CommandDispatcher;
+
 import me.lucko.spark.fabric.plugin.FabricClientSparkPlugin;
 import me.lucko.spark.fabric.plugin.FabricServerSparkPlugin;
 
 import net.fabricmc.api.ModInitializer;
+import net.fabricmc.fabric.api.command.v1.CommandRegistrationCallback;
 import net.fabricmc.fabric.api.event.lifecycle.v1.ServerLifecycleEvents;
 import net.fabricmc.loader.api.FabricLoader;
 import net.fabricmc.loader.api.ModContainer;
 import net.minecraft.client.MinecraftClient;
+import net.minecraft.server.MinecraftServer;
+import net.minecraft.server.command.ServerCommandSource;
 
 import java.nio.file.Path;
 import java.util.Objects;
@@ -38,6 +43,8 @@ public class FabricSparkMod implements ModInitializer {
     private ModContainer container;
     private Path configDirectory;
 
+    private FabricServerSparkPlugin activeServerPlugin = null;
+
     @Override
     public void onInitialize() {
         FabricSparkMod.mod = this;
@@ -47,8 +54,12 @@ public class FabricSparkMod implements ModInitializer {
                 .orElseThrow(() -> new IllegalStateException("Unable to get container for spark"));
         this.configDirectory = loader.getConfigDir().resolve("spark");
 
-        // load hooks
-        ServerLifecycleEvents.SERVER_STARTING.register(server -> FabricServerSparkPlugin.register(this, server));
+        // lifecycle hooks
+        ServerLifecycleEvents.SERVER_STARTING.register(this::initializeServer);
+        ServerLifecycleEvents.SERVER_STOPPING.register(this::onServerStopping);
+
+        // events to propagate to active server plugin
+        CommandRegistrationCallback.EVENT.register(this::onCommandRegister);
     }
 
     // called be entrypoint defined in fabric.mod.json
@@ -57,6 +68,23 @@ public class FabricSparkMod implements ModInitializer {
         FabricClientSparkPlugin.register(FabricSparkMod.mod, MinecraftClient.getInstance());
     }
 
+    public void initializeServer(MinecraftServer server) {
+        this.activeServerPlugin = FabricServerSparkPlugin.register(this, server);
+    }
+
+    public void onServerStopping(MinecraftServer stoppingServer) {
+        if (this.activeServerPlugin != null) {
+            this.activeServerPlugin.disable();
+            this.activeServerPlugin = null;
+        }
+    }
+
+    public void onCommandRegister(CommandDispatcher<ServerCommandSource> dispatcher, boolean isDedicated) {
+        if (this.activeServerPlugin != null) {
+            this.activeServerPlugin.registerCommands(dispatcher);
+        }
+    }
+
     public String getVersion() {
         return this.container.getMetadata().getVersion().getFriendlyString();
     }