refactor sanity and insanity commands; add max sanity handling and improve syncing logic

src/main/java/net/fretux/mindmotion/command/SanityCommand.java

@@ -15,21 +15,22 @@ public class SanityCommand {
         dispatcher.register(Commands.literal("sanity")
             .requires(source -> source.hasPermission(2))
             .then(Commands.literal("set")
-                .then(Commands.argument("sanity", FloatArgumentType.floatArg(0, 80))
+                .then(Commands.argument("sanity", FloatArgumentType.floatArg(0))
-                    .then(Commands.argument("insanity", FloatArgumentType.floatArg(0, 80))
                     .executes(ctx -> {
-                            ServerPlayer player = ctx.getSource().getPlayer();
+                        ServerPlayer player = ctx.getSource().getPlayerOrException();
-                            float sanity = FloatArgumentType.getFloat(ctx, "sanity");
+                        float requestedSanity = FloatArgumentType.getFloat(ctx, "sanity");
-                            float insanity = FloatArgumentType.getFloat(ctx, "insanity");
+                        final float[] appliedSanity = new float[] {0f};
-                            assert player != null;
                         player.getCapability(PlayerCapabilityProvider.SANITY).ifPresent(cap -> {
-                                cap.setSanity(sanity);
+                            float clampedSanity = Math.min(cap.getMaxSanity(), requestedSanity);
-                                cap.setInsanity(insanity);
+                            if (clampedSanity > 0f && cap.getInsanity() > 0f) {
+                                cap.setInsanity(0f);
+                            }
+                            cap.setSanity(clampedSanity);
+                            appliedSanity[0] = cap.getSanity();
                         });
 
                         ctx.getSource().sendSuccess(() ->
-                                net.minecraft.network.chat.Component.literal(
+                            net.minecraft.network.chat.Component.literal("Set sanity = " + appliedSanity[0]),
-                                        "Set sanity = " + sanity + ", insanity = " + insanity),
                             true
                         );
 
@@ -37,6 +38,32 @@ public class SanityCommand {
                     })
                 )
             )
+        );
+        dispatcher.register(Commands.literal("insanity")
+            .requires(source -> source.hasPermission(2))
+            .then(Commands.literal("set")
+                .then(Commands.argument("insanity", FloatArgumentType.floatArg(0))
+                    .executes(ctx -> {
+                        ServerPlayer player = ctx.getSource().getPlayerOrException();
+                        float requestedInsanity = FloatArgumentType.getFloat(ctx, "insanity");
+                        final float[] appliedInsanity = new float[] {0f};
+                        player.getCapability(PlayerCapabilityProvider.SANITY).ifPresent(cap -> {
+                            float clampedInsanity = Math.min(cap.getMaxSanity(), requestedInsanity);
+                            if (clampedInsanity > 0f && cap.getSanity() > 0f) {
+                                cap.setSanity(0f);
+                            }
+                            cap.setInsanity(clampedInsanity);
+                            appliedInsanity[0] = cap.getInsanity();
+                        });
+
+                        ctx.getSource().sendSuccess(() ->
+                            net.minecraft.network.chat.Component.literal("Set insanity = " + appliedInsanity[0]),
+                            true
+                        );
+
+                        return 1;
+                    })
+                )
             )
         );
     }

src/main/java/net/fretux/mindmotion/event/CommonEvents.java

@@ -1,6 +1,7 @@
 package net.fretux.mindmotion.event;
 
 import net.fretux.mindmotion.player.PlayerCapabilityProvider;
+import net.fretux.mindmotion.player.SanityCapability;
 import net.minecraftforge.event.entity.living.LivingHurtEvent;
 import net.minecraftforge.event.entity.player.PlayerEvent;
 import net.minecraftforge.eventbus.api.SubscribeEvent;
@@ -21,8 +22,17 @@ public class CommonEvents {
         event.getEntity().getCapability(PlayerCapabilityProvider.SANITY).ifPresent(newCap ->
             event.getOriginal().getCapability(PlayerCapabilityProvider.SANITY)
                     .ifPresent(oldCap -> {
+                        if (newCap instanceof SanityCapability newImpl && oldCap instanceof SanityCapability oldImpl) {
+                            newImpl.setBaseMaxSanity(oldImpl.getBaseMaxSanity());
+                            newImpl.setBonusMaxSanity(oldImpl.getBonusMaxSanity());
+                        }
+                        if (event.isWasDeath()) {
+                            newCap.setInsanity(0f);
+                            newCap.setSanity(newCap.getMaxSanity());
+                        } else {
                             newCap.setSanity(oldCap.getSanity());
                             newCap.setInsanity(oldCap.getInsanity());
+                        }
                     })
         );
         event.getEntity().getCapability(PlayerCapabilityProvider.TEMPO).ifPresent(newCap ->

src/main/java/net/fretux/mindmotion/player/PlayerCapabilityProvider.java

@@ -40,6 +40,8 @@ public class PlayerCapabilityProvider {
         @Override
         public CompoundTag serializeNBT() {
             CompoundTag tag = new CompoundTag();
+            tag.putFloat("BaseMaxSanity", instance.getBaseMaxSanity());
+            tag.putFloat("BonusMaxSanity", instance.getBonusMaxSanity());
             tag.putFloat("Sanity", instance.getSanity());
             tag.putFloat("Insanity", instance.getInsanity());
             return tag;
@@ -47,6 +49,12 @@ public class PlayerCapabilityProvider {
 
         @Override
         public void deserializeNBT(CompoundTag nbt) {
+            if (nbt.contains("BaseMaxSanity")) {
+                instance.setBaseMaxSanity(nbt.getFloat("BaseMaxSanity"));
+            }
+            if (nbt.contains("BonusMaxSanity")) {
+                instance.setBonusMaxSanity(nbt.getFloat("BonusMaxSanity"));
+            }
             instance.setSanity(nbt.getFloat("Sanity"));
             instance.setInsanity(nbt.getFloat("Insanity"));
         }

src/main/java/net/fretux/mindmotion/player/SanityCapability.java

@@ -57,6 +57,10 @@ public class SanityCapability implements ISanity {
         setInsanity(this.insanity);
     }
 
+    public float getBaseMaxSanity() {
+        return baseMaxSanity;
+    }
+
     public void setBonusMaxSanity(float bonusMaxSanity) {
         this.bonusMaxSanity = Math.max(0f, bonusMaxSanity);
         clampToMax();