From d85f83bfa8023559636ca5b479a925f2fd7a44ff Mon Sep 17 00:00:00 2001
From: Kai <kai.zahn@yahoo.de>
Date: Sat, 23 Nov 2024 12:58:43 +0100
Subject: [PATCH] assign stage when the stored one is no longer present

---
 .../common/capability/DragonStateHandler.java | 22 +++++++++----------
 1 file changed, 11 insertions(+), 11 deletions(-)

diff --git a/src/main/java/by/dragonsurvivalteam/dragonsurvival/common/capability/DragonStateHandler.java b/src/main/java/by/dragonsurvivalteam/dragonsurvival/common/capability/DragonStateHandler.java
index abf6ae61fa..413363c5af 100644
--- a/src/main/java/by/dragonsurvivalteam/dragonsurvival/common/capability/DragonStateHandler.java
+++ b/src/main/java/by/dragonsurvivalteam/dragonsurvival/common/capability/DragonStateHandler.java
@@ -509,15 +509,19 @@ public void deserializeNBT(HolderLookup.Provider provider, CompoundTag tag, bool
                             () -> DragonSurvival.LOGGER.warn("Cannot set dragon body [{}] while deserializing NBT of [{}] due to the dragon body not existing", storedDragonBody, tag));
         }
 
+        double size = tag.getDouble(SIZE);
+        Holder<DragonStage> dragonStage = null;
+
         String storedDragonStage = tag.getString(DRAGON_STAGE);
 
         if (!storedDragonStage.isEmpty()) {
-            Holder<DragonStage> dragonStage = loadStage(provider, storedDragonStage);
+            Holder<DragonStage> loadedDragonStage = loadStage(provider, storedDragonStage);
 
-            if (dragonStage != null) {
-                this.dragonStage = dragonStage;
-            } else {
-                DragonSurvival.LOGGER.warn("Cannot set dragon stage [{}] while deserializing NBT of [{}] due to the dragon stage not existing", storedDragonStage, tag);
+            if (loadedDragonStage != null) {
+                dragonStage = loadedDragonStage;
+            } else if (size != NO_SIZE) {
+                dragonStage = DragonStage.get(provider, size);
+                DragonSurvival.LOGGER.warn("Cannot set dragon stage [{}] while deserializing NBT of [{}] due to the dragon stage not existing - setting [{}] as fallback", storedDragonStage, tag, dragonStage);
             }
         }
 
@@ -525,7 +529,7 @@ public void deserializeNBT(HolderLookup.Provider provider, CompoundTag tag, bool
 
         if (!storedPreviousDragonStage.isEmpty()) {
             // This is not interesting enough to the user to log a warning
-            provider.holder(DragonStages.key(ResourceLocation.parse(storedPreviousDragonStage))).ifPresent(dragonStage -> this.previousStage = dragonStage);
+            provider.holder(DragonStages.key(ResourceLocation.parse(storedPreviousDragonStage))).ifPresent(previousStage -> this.previousStage = previousStage);
         }
 
         if (dragonType != null) {
@@ -545,11 +549,7 @@ public void deserializeNBT(HolderLookup.Provider provider, CompoundTag tag, bool
             getMovementData().spinCooldown = tag.getInt("spinCooldown");
             getMovementData().spinAttack = tag.getInt("spinAttack");
 
-            if (dragonStage == null) {
-                setClientSize(tag.getDouble(SIZE));
-            } else {
-                setClientSize(dragonStage, tag.getDouble(SIZE));
-            }
+            setClientSize(dragonStage, size);
 
             setDestructionEnabled(tag.getBoolean("destructionEnabled"));
             isGrowing = !tag.contains(IS_GROWING) || tag.getBoolean(IS_GROWING);