Added stage cleanup feature

src/main/java/com/github/hhhzzzsss/songplayer/CommandProcessor.java

@@ -9,12 +9,17 @@ import com.github.hhhzzzsss.songplayer.song.Playlist;
 import com.github.hhhzzzsss.songplayer.song.Song;
 import com.mojang.brigadier.suggestion.Suggestions;
 import com.mojang.brigadier.suggestion.SuggestionsBuilder;
+import net.minecraft.block.BlockState;
 import net.minecraft.command.CommandSource;
 import net.minecraft.component.DataComponentTypes;
 import net.minecraft.component.type.NbtComponent;
 import net.minecraft.item.ItemStack;
 import net.minecraft.nbt.NbtCompound;
+import net.minecraft.state.property.Property;
+import net.minecraft.text.*;
+import net.minecraft.util.Formatting;
 import net.minecraft.util.Hand;
+import net.minecraft.util.hit.BlockHitResult;
 import net.minecraft.world.GameMode;
 
 import java.io.IOException;
@@ -52,9 +57,12 @@ public class CommandProcessor {
 		commands.add(new toggleFakePlayerCommand());
 		commands.add(new setStageTypeCommand());
 		commands.add(new toggleMovementCommand());
+		commands.add(new toggleAutoCleanup());
+		commands.add(new cleanupLastStageCommand());
 		commands.add(new announcementCommand());
 		commands.add(new songItemCommand());
 		commands.add(new testSongCommand());
+		commands.add(new testBlockStateCommand());
 
 		for (Command command : commands) {
 			commandMap.put(command.getName().toLowerCase(Locale.ROOT), command);
@@ -236,16 +244,21 @@ public class CommandProcessor {
 			return "Stops playing";
 		}
 		public boolean processCommand(String args) {
-			if (SongHandler.getInstance().currentSong == null && SongHandler.getInstance().songQueue.isEmpty()) {
+			if (SongHandler.getInstance().isIdle()) {
 				SongPlayer.addChatMessage("§6No song is currently playing");
 				return true;
 			}
 			if (args.length() == 0) {
-				if (SongHandler.getInstance().stage != null) {
+				if (SongHandler.getInstance().cleaningUp) {
-					SongHandler.getInstance().stage.movePlayerToStagePosition();
+					SongHandler.getInstance().restoreStateAndReset();
+					SongPlayer.addChatMessage("§6Stopped cleanup");
+				} else if (Config.getConfig().autoCleanup && SongHandler.getInstance().originalBlocks.size() != 0) {
+					SongHandler.getInstance().partionResetAndCleanup();
+					SongPlayer.addChatMessage("§6Stopped playing and switched to cleanup");
+				} else {
+					SongHandler.getInstance().restoreStateAndReset();
+					SongPlayer.addChatMessage("§6Stopped playing");
 				}
-				SongHandler.getInstance().restoreStateAndCleanUp();
-				SongPlayer.addChatMessage("§6Stopped playing");
 				return true;
 			}
 			else {
@@ -935,6 +948,92 @@ public class CommandProcessor {
 		}
 	}
 
+	private static class toggleAutoCleanup extends Command {
+		public String getName() {
+			return "toggleAutoCleanup";
+		}
+		public String[] getAliases() {
+			return new String[]{"autoCleanup"};
+		}
+		public String[] getSyntax() {
+			return new String[0];
+		}
+		public String getDescription() {
+			return "Toggles whether you automatically clean up your stage and restore the original blocks after playing";
+		}
+		public boolean processCommand(String args) {
+			if (args.length() == 0) {
+				Config.getConfig().autoCleanup = !Config.getConfig().autoCleanup;
+				if (Config.getConfig().autoCleanup) {
+					SongPlayer.addChatMessage("§6Enabled automatic cleanup");
+				}
+				else {
+					SongPlayer.addChatMessage("§6Disabled automatic cleanup");
+				}
+				Config.saveConfigWithErrorHandling();
+				return true;
+			}
+			else {
+				return false;
+			}
+		}
+	}
+
+	private static class cleanupLastStageCommand extends Command {
+		public String getName() {
+			return "cleanupLastStage";
+		}
+		public String[] getAliases() {
+			return new String[]{};
+		}
+		public String[] getSyntax() {
+			return new String[0];
+		}
+		public String getDescription() {
+			return "Cleans up your most recent stage and restores the original blocks";
+		}
+		public boolean processCommand(String args) {
+			if (args.length() == 0) {
+				Stage lastStage = SongHandler.getInstance().lastStage;
+				if (!SongHandler.getInstance().isIdle()) {
+					SongPlayer.addChatMessage("§cYou cannot start cleanup if you are in the middle of another action");
+					return true;
+				}
+				if (lastStage == null || SongHandler.getInstance().originalBlocks.size() == 0) {
+					SongPlayer.addChatMessage("§6There is nothing to clean up");
+					return true;
+				}
+				if (SongPlayer.MC.player.getPos().squaredDistanceTo(lastStage.getOriginBottomCenter()) > 3*3) {
+					System.out.println(SongPlayer.MC.player.getPos().squaredDistanceTo(lastStage.getOriginBottomCenter()));
+					String coordStr = String.format(
+							"%d %d %d",
+							lastStage.position.getX(), lastStage.position.getY(), lastStage.position.getZ()
+					);
+					SongPlayer.addChatMessage("§6You must be within §33 §6blocks of the center of your stage to start cleanup.");
+					MutableText coordText = Util.joinTexts(null,
+							Text.literal("This is at ").setStyle(Style.EMPTY.withColor(Formatting.GOLD)),
+							Text.literal(coordStr).setStyle(
+									Style.EMPTY
+											.withColor(Formatting.DARK_AQUA)
+											.withUnderline(true)
+											.withClickEvent(new ClickEvent(ClickEvent.Action.COPY_TO_CLIPBOARD, coordStr))
+											.withHoverEvent(new HoverEvent(HoverEvent.Action.SHOW_TEXT, Text.literal("Copy \"" + coordStr + "\"")))
+							),
+							Text.literal(" (click to copy)").setStyle(Style.EMPTY.withColor(Formatting.GOLD))
+					);
+					SongPlayer.addChatMessage(coordText);
+					return true;
+				}
+
+				SongHandler.getInstance().startCleanup();
+				return true;
+			}
+			else {
+				return false;
+			}
+		}
+	}
+
 	private static class announcementCommand extends Command {
 		public String getName() {
 			return "announcement";
@@ -1093,6 +1192,39 @@ public class CommandProcessor {
 		}
 	}
 
+	private static class testBlockStateCommand extends Command {
+		public String getName() {
+			return "testBlockState";
+		}
+		public String[] getSyntax() {
+			return new String[0];
+		}
+		public String getDescription() {
+			return "for dev purposes";
+		}
+		public boolean processCommand(String args) {
+			if (args.length() == 0) {
+				if (MC.crosshairTarget instanceof BlockHitResult) {
+					BlockHitResult hitResult = (BlockHitResult) MC.crosshairTarget;
+					BlockState bs = MC.world.getBlockState(hitResult.getBlockPos());
+					ItemStack stack = new ItemStack(bs.getBlock());
+					SongPlayer.addChatMessage(stack.toString());
+					for (Map.Entry<Property<?>, Comparable<?>> entry : bs.getEntries().entrySet()) {
+						Property<?> property = entry.getKey();
+						Comparable<?> value = entry.getValue();
+//						SongPlayer.addChatMessage(net.minecraft.util.Util.getValueAsString(property, comparable));
+						System.out.println(property.getClass());
+						System.out.println(value.getClass());
+					}
+				}
+				return true;
+			}
+			else {
+				return false;
+			}
+		}
+	}
+
 	public static CompletableFuture<Suggestions> handleSuggestions(String text, SuggestionsBuilder suggestionsBuilder) {
 		if (!text.contains(" ")) {
 			List<String> names = commandCompletions

src/main/java/com/github/hhhzzzsss/songplayer/Config.java

@@ -26,6 +26,7 @@ public class Config {
     public boolean rotate = false;
     public boolean doAnnouncement = false;
     public String announcementMessage = "&6Now playing: &3[name]";
+    public boolean autoCleanup = false;
 
     public static Config getConfig() {
         if (config == null) {

src/main/java/com/github/hhhzzzsss/songplayer/FakePlayerEntity.java

@@ -49,9 +49,9 @@ public class FakePlayerEntity extends OtherClientPlayerEntity {
 	}
 	
 	public void copyStagePosAndPlayerLook() {
-		Stage stage = SongHandler.getInstance().stage;
+		Stage lastStage = SongHandler.getInstance().lastStage;
-		if (stage != null) {
+		if (lastStage != null) {
-			refreshPositionAndAngles(stage.position.getX()+0.5, stage.position.getY(), stage.position.getZ()+0.5, player.getYaw(), player.getPitch());
+			refreshPositionAndAngles(lastStage.position.getX()+0.5, lastStage.position.getY(), lastStage.position.getZ()+0.5, player.getYaw(), player.getPitch());
 			headYaw = player.headYaw;
 		}
 		else {

src/main/java/com/github/hhhzzzsss/songplayer/Util.java

@@ -6,10 +6,8 @@ import net.minecraft.command.CommandSource;
 import net.minecraft.component.DataComponentTypes;
 import net.minecraft.component.type.LoreComponent;
 import net.minecraft.item.ItemStack;
-import net.minecraft.nbt.NbtList;
-import net.minecraft.nbt.NbtString;
-import net.minecraft.text.PlainTextContent;
 import net.minecraft.text.MutableText;
+import net.minecraft.text.PlainTextContent;
 import net.minecraft.text.Style;
 import net.minecraft.text.Text;
 
@@ -230,4 +228,14 @@ public class Util {
     public static void setItemLore(ItemStack stack, Text... loreLines) {
         stack.set(DataComponentTypes.LORE, new LoreComponent(List.of(loreLines)));
     }
+
+    public static MutableText joinTexts(MutableText base, Text... children) {
+        if (base == null) {
+            base = Text.empty();
+        }
+        for (Text child : children) {
+            base.append(child);
+        }
+        return base;
+    }
 }

src/main/java/com/github/hhhzzzsss/songplayer/mixin/ClientCommonNetworkHandlerMixin.java

@@ -27,11 +27,11 @@ public class ClientCommonNetworkHandlerMixin {
 
     @Inject(at = @At("HEAD"), method = "sendPacket(Lnet/minecraft/network/packet/Packet;)V", cancellable = true)
     private void onSendPacket(Packet<?> packet, CallbackInfo ci) {
-        Stage stage = SongHandler.getInstance().stage;
+        Stage lastStage = SongHandler.getInstance().lastStage;
 
-        if (stage != null && packet instanceof PlayerMoveC2SPacket) {
+        if (!SongHandler.getInstance().isIdle() && lastStage != null && packet instanceof PlayerMoveC2SPacket) {
             if (!Config.getConfig().rotate) {
-                connection.send(new PlayerMoveC2SPacket.Full(stage.position.getX() + 0.5, stage.position.getY(), stage.position.getZ() + 0.5, SongPlayer.MC.player.getYaw(), SongPlayer.MC.player.getPitch(), true));
+                connection.send(new PlayerMoveC2SPacket.Full(lastStage.position.getX() + 0.5, lastStage.position.getY(), lastStage.position.getZ() + 0.5, SongPlayer.MC.player.getYaw(), SongPlayer.MC.player.getPitch(), true));
                 if (SongPlayer.fakePlayer != null) {
                     SongPlayer.fakePlayer.copyStagePosAndPlayerLook();
                 }

src/main/java/com/github/hhhzzzsss/songplayer/mixin/ClientPlayNetworkHandlerMixin.java

@@ -24,27 +24,54 @@ public class ClientPlayNetworkHandlerMixin {
 	
 	@Inject(at = @At("TAIL"), method = "onGameJoin(Lnet/minecraft/network/packet/s2c/play/GameJoinS2CPacket;)V")
 	public void onOnGameJoin(GameJoinS2CPacket packet, CallbackInfo ci) {
-		SongHandler.getInstance().cleanup();
+		SongHandler.getInstance().reset();
 	}
 
 	@Inject(at = @At("TAIL"), method = "onPlayerRespawn(Lnet/minecraft/network/packet/s2c/play/PlayerRespawnS2CPacket;)V")
 	public void onOnPlayerRespawn(PlayerRespawnS2CPacket packet, CallbackInfo ci) {
-		SongHandler.getInstance().cleanup();
+		SongHandler.getInstance().reset();
 	}
 
 	@Inject(at = @At("TAIL"), method = "onPlayerPositionLook(Lnet/minecraft/network/packet/s2c/play/PlayerPositionLookS2CPacket;)V")
 	public void onOnPlayerPositionLook(PlayerPositionLookS2CPacket packet, CallbackInfo ci) {
-		Stage stage = SongHandler.getInstance().stage;
+		Stage lastStage = SongHandler.getInstance().lastStage;
-		if (!SongHandler.getInstance().isIdle() && stage != null && Vec3d.ofBottomCenter(stage.position).squaredDistanceTo(SongPlayer.MC.player.getPos()) > 3*3) {
+		if (!SongHandler.getInstance().isIdle() && lastStage != null && lastStage.getOriginBottomCenter().squaredDistanceTo(SongPlayer.MC.player.getPos()) > 3*3) {
-			SongPlayer.addChatMessage("§6Stopped playing because the server moved the player too far from the stage!");
+			Vec3d stageOriginBottomCenter = lastStage.getOriginBottomCenter();
-			SongHandler.getInstance().cleanup();
+			boolean xrel = packet.getFlags().contains(PositionFlag.X);
+			boolean yrel = packet.getFlags().contains(PositionFlag.Y);
+			boolean zrel = packet.getFlags().contains(PositionFlag.Z);
+			double dx = 0.0;
+			double dy = 0.0;
+			double dz = 0.0;
+			if (xrel) {
+				dx = packet.getX();
+			} else {
+				dx = SongPlayer.MC.player.getX() - stageOriginBottomCenter.getX();
+			}
+			if (yrel) {
+				dy = packet.getY();
+			} else {
+				dy = SongPlayer.MC.player.getY() - stageOriginBottomCenter.getY();
+			}
+			if (zrel) {
+				dz = packet.getZ();
+			} else {
+				dz = SongPlayer.MC.player.getZ() - stageOriginBottomCenter.getZ();
+			}
+			double dist = dx*dx + dy*dy + dz*dz;
+			if (dist > 3.0) {
+				SongPlayer.addChatMessage("§6Stopped playing/building because the server moved the player too far from the stage!");
+				SongHandler.getInstance().reset();
+			} else {
+				lastStage.movePlayerToStagePosition();
+			}
 		}
 	}
 
 	@Inject(at = @At("TAIL"), method = "onPlayerAbilities(Lnet/minecraft/network/packet/s2c/play/PlayerAbilitiesS2CPacket;)V")
 	public void onOnPlayerAbilities(PlayerAbilitiesS2CPacket packet, CallbackInfo ci) {
 		SongHandler handler = SongHandler.getInstance();
-		if (handler.currentSong != null || handler.currentPlaylist != null || handler.songQueue.size() > 0) {
+		if (handler.wasFlying) {
 			SongPlayer.MC.player.getAbilities().flying = handler.wasFlying;
 		}
 	}

src/main/java/com/github/hhhzzzsss/songplayer/playing/SongHandler.java

@@ -6,7 +6,7 @@ import com.github.hhhzzzsss.songplayer.SongPlayer;
 import com.github.hhhzzzsss.songplayer.Util;
 import com.github.hhhzzzsss.songplayer.mixin.ClientPlayerInteractionManagerAccessor;
 import com.github.hhhzzzsss.songplayer.song.*;
-import net.minecraft.block.Block;
+import net.minecraft.block.*;
 import net.minecraft.client.world.ClientWorld;
 import net.minecraft.component.DataComponentTypes;
 import net.minecraft.component.type.BlockStateComponent;
@@ -14,8 +14,8 @@ import net.minecraft.entity.player.PlayerInventory;
 import net.minecraft.item.ItemStack;
 import net.minecraft.item.Items;
 import net.minecraft.network.packet.c2s.play.PlayerMoveC2SPacket;
-import net.minecraft.text.MutableText;
+import net.minecraft.state.property.Property;
-import net.minecraft.text.Text;
+import net.minecraft.text.*;
 import net.minecraft.util.Formatting;
 import net.minecraft.util.Hand;
 import net.minecraft.util.hit.BlockHitResult;
@@ -27,8 +27,8 @@ import net.minecraft.world.GameMode;
 
 import java.io.IOException;
 import java.nio.file.Path;
-import java.util.LinkedList;
+import java.util.*;
-import java.util.Map;
+import java.util.stream.Collectors;
 
 public class SongHandler {
     private static SongHandler instance = null;
@@ -44,8 +44,12 @@ public class SongHandler {
     public LinkedList<Song> songQueue = new LinkedList<>();
     public Song currentSong = null;
     public Playlist currentPlaylist = null;
-    public Stage stage = null;
+    public Stage stage = null; // Only exists when playing
+    public Stage lastStage = null; // Stays around even after playing
+    public HashMap<BlockPos, BlockState> originalBlocks = new HashMap<>();
     public boolean building = false;
+    public boolean cleaningUp = false;
+    public boolean dirty = false;
 
     public boolean wasFlying = false;
     public GameMode originalGamemode = GameMode.CREATIVE;
@@ -53,66 +57,54 @@ public class SongHandler {
     boolean playlistChecked = false;
 
     public void onUpdate(boolean tick) {
-        // Check current playlist and load song from it if necessary
+        if (!cleaningUp) {
-        if (currentSong == null && currentPlaylist != null && currentPlaylist.loaded) {
+            // Check current playlist and load song from it if necessary
-            if (!playlistChecked) {
+            if (currentSong == null && currentPlaylist != null && currentPlaylist.loaded) {
-                playlistChecked = true;
+                if (!playlistChecked) {
-                if (currentPlaylist.songsFailedToLoad.size() > 0) {
+                    playlistChecked = true;
-                    SongPlayer.addChatMessage("§cFailed to load the following songs from the playlist: §4" + String.join(" ", currentPlaylist.songsFailedToLoad));
+                    if (currentPlaylist.songsFailedToLoad.size() > 0) {
+                        SongPlayer.addChatMessage("§cFailed to load the following songs from the playlist: §4" + String.join(" ", currentPlaylist.songsFailedToLoad));
+                    }
                 }
-            }
+                Song nextSong = currentPlaylist.getNext();
-            Song nextSong = currentPlaylist.getNext();
+                if (currentPlaylist.songs.size() == 0) {
-            if (currentPlaylist.songs.size() == 0) {
+                    SongPlayer.addChatMessage("§cPlaylist has no playable songs");
-                SongPlayer.addChatMessage("§cPlaylist has no playable songs");
+                    currentPlaylist = null;
-                currentPlaylist = null;
+                } else if (nextSong == null) {
-            }
+                    SongPlayer.addChatMessage("§6Playlist has finished playing");
-            else if (nextSong == null) {
+                    currentPlaylist = null;
-                SongPlayer.addChatMessage("§6Playlist has finished playing");
-                currentPlaylist = null;
-            }
-            else {
-                nextSong.reset();
-                setSong(nextSong);
-            }
-        }
-
-        // Check queue and load song from it if necessary
-        if (currentSong == null && currentPlaylist == null && songQueue.size() > 0) {
-            setSong(songQueue.poll());
-        }
-
-        // Check if loader thread is finished and handle accordingly
-        if (loaderThread != null && !loaderThread.isAlive()) {
-            if (loaderThread.exception != null) {
-                SongPlayer.addChatMessage("§cFailed to load song: §4" + loaderThread.exception.getMessage());
-            } else {
-                if (currentSong == null) {
-                    setSong(loaderThread.song);
                 } else {
-                    queueSong(loaderThread.song);
+                    nextSong.reset();
+                    setSong(nextSong);
                 }
             }
-            loaderThread = null;
+
+            // Check queue and load song from it if necessary
+            if (currentSong == null && currentPlaylist == null && songQueue.size() > 0) {
+                setSong(songQueue.poll());
+            }
+
+            // Check if loader thread is finished and handle accordingly
+            if (loaderThread != null && !loaderThread.isAlive()) {
+                if (loaderThread.exception != null) {
+                    SongPlayer.addChatMessage("§cFailed to load song: §4" + loaderThread.exception.getMessage());
+                } else {
+                    if (currentSong == null) {
+                        setSong(loaderThread.song);
+                    } else {
+                        queueSong(loaderThread.song);
+                    }
+                }
+                loaderThread = null;
+            }
         }
 
         // Run cached command if timeout reached
         checkCommandCache();
 
-        // Check if no song is playing and, if necessary, handle cleanup
+        // If either playing or doing cleanup
-        if (currentSong == null) {
+        if (cleaningUp || currentSong != null) {
-            if (stage != null || SongPlayer.fakePlayer != null) {
+            // Handle creating/removing fake player depending on settings
-                restoreStateAndCleanUp();
-            }
-            else {
-                originalGamemode = SongPlayer.MC.interactionManager.getCurrentGameMode();
-            }
-        }
-        // Otherwise, handle song playing
-        else {
-            if (stage == null) {
-                stage = new Stage();
-                stage.movePlayerToStagePosition();
-            }
             if (Config.getConfig().showFakePlayer && SongPlayer.fakePlayer == null) {
                 SongPlayer.fakePlayer = new FakePlayerEntity();
                 SongPlayer.fakePlayer.copyStagePosAndPlayerLook();
@@ -124,9 +116,31 @@ public class SongHandler {
                 SongPlayer.fakePlayer.getInventory().clone(SongPlayer.MC.player.getInventory());
             }
 
+            // Allow flying
             SongPlayer.MC.player.getAbilities().allowFlying = true;
             wasFlying = SongPlayer.MC.player.getAbilities().flying;
+        }
 
+        // Check if doing cleanup
+        if (cleaningUp) {
+            if (tick) {
+                // Allow flying while doing cleanup
+                SongPlayer.MC.player.getAbilities().allowFlying = true;
+                wasFlying = SongPlayer.MC.player.getAbilities().flying;
+
+                handleCleanup();
+            }
+        }
+        // Check if song is playing
+        else if (currentSong != null) {
+            // This should never happen, but I left this check in just in case.
+            if (stage == null) {
+                SongPlayer.addChatMessage("§cStage is null! This should not happen!");
+                reset();
+                return;
+            }
+
+            // Run building or playing tick depending on state
             if (building) {
                 if (tick) {
                     handleBuilding();
@@ -135,6 +149,20 @@ public class SongHandler {
                 handlePlaying(tick);
             }
         }
+        // Otherwise, handle cleanup if necessary
+        else {
+            if (dirty) {
+                if (Config.getConfig().autoCleanup && originalBlocks.size() != 0) {
+                    partionResetAndCleanup();
+                } else {
+                    restoreStateAndReset();
+                }
+            }
+            else {
+                // When doing nothing else, record original gamemode
+                originalGamemode = SongPlayer.MC.interactionManager.getCurrentGameMode();
+            }
+        }
     }
 
     public void loadSong(String location) {
@@ -167,23 +195,18 @@ public class SongHandler {
         }
     }
 
+    // Sets currentSong and sets everything up for building
     public void setSong(Song song) {
+        dirty = true;
         currentSong = song;
         building = true;
         setCreativeIfNeeded();
         if (Config.getConfig().doAnnouncement) {
             sendMessage(Config.getConfig().announcementMessage.replaceAll("\\[name\\]", song.name));
         }
-        if (stage == null) {
+        prepareStage();
-            stage = new Stage();
-            stage.movePlayerToStagePosition();
-        }
-        else {
-            stage.sendMovementPacketToStagePosition();
-        }
         getAndSaveBuildSlot();
         SongPlayer.addChatMessage("§6Building noteblocks");
-
     }
 
     private void queueSong(Song song) {
@@ -213,6 +236,14 @@ public class SongHandler {
         }
     }
 
+    public void startCleanup() {
+        dirty = true;
+        cleaningUp = true;
+        setCreativeIfNeeded();
+        getAndSaveBuildSlot();
+        lastStage.sendMovementPacketToStagePosition();
+    }
+
     // Runs every tick
     private int buildStartDelay = 0;
     private int buildEndDelay = 0;
@@ -240,6 +271,7 @@ public class SongHandler {
                 return;
             } else {
                 stage.checkBuildStatus(currentSong);
+                recordStageBlocks();
                 stage.sendMovementPacketToStagePosition();
             }
         }
@@ -310,6 +342,7 @@ public class SongHandler {
         if (tick) {
             if (stage.hasBreakingModification()) {
                 stage.checkBuildStatus(currentSong);
+                recordStageBlocks();
             }
             if (!stage.nothingToBuild()) { // Switch to building
                 building = true;
@@ -350,8 +383,7 @@ public class SongHandler {
             currentSong = null;
         }
     }
-
+    private void setPlayProgressDisplay() {
-    public void setPlayProgressDisplay() {
         long currentTime = Math.min(currentSong.time, currentSong.length);
         long totalTime = currentSong.length;
         MutableText songText = Text.empty()
@@ -382,18 +414,201 @@ public class SongHandler {
         ProgressDisplay.getInstance().setText(songText, playlistText);
     }
 
-    public void cleanup() {
+    // Runs every tick
+    private int cleanupTotalBlocksToPlace = 0;
+    private LinkedList<BlockPos> cleanupBreakList = new LinkedList<>();
+    private LinkedList<BlockPos> cleanupPlaceList = new LinkedList<>();
+    private ArrayList<BlockPos> cleanupUnplaceableBlocks = new ArrayList<>();
+    private void handleCleanup() {
+        setCleanupProgressDisplay();
+
+        if (buildStartDelay > 0) {
+            buildStartDelay--;
+            return;
+        }
+        if (buildCooldown > 0) {
+            buildCooldown--;
+            return;
+        }
+        ClientWorld world = SongPlayer.MC.world;
+        if (SongPlayer.MC.interactionManager.getCurrentGameMode() != GameMode.CREATIVE) {
+            return;
+        }
+
+        if (cleanupBreakList.isEmpty() && cleanupPlaceList.isEmpty()) {
+            if (buildEndDelay > 0) {
+                buildEndDelay--;
+                return;
+            } else {
+                checkCleanupStatus();
+                lastStage.sendMovementPacketToStagePosition();
+            }
+        }
+
+        if (!cleanupBreakList.isEmpty()) {
+            for (int i=0; i<5; i++) {
+                if (cleanupBreakList.isEmpty()) break;
+                BlockPos bp = cleanupBreakList.poll();
+                attackBlock(bp);
+            }
+            buildEndDelay = 20;
+        } else if (!cleanupPlaceList.isEmpty()) {
+            BlockPos bp = cleanupPlaceList.pollFirst();
+            BlockState actualBlockState = world.getBlockState(bp);
+            BlockState desiredBlockState = originalBlocks.get(bp);
+            if (actualBlockState != desiredBlockState) {
+                holdBlock(desiredBlockState, buildSlot);
+                if (!actualBlockState.isAir() && !actualBlockState.isLiquid()) {
+                    attackBlock(bp);
+                }
+                placeBlock(bp);
+            }
+            buildCooldown = 0; // No cooldown, so it places a block every tick
+            buildEndDelay = 20;
+        } else {
+            originalBlocks.clear();
+            cleaningUp = false;
+            SongPlayer.addChatMessage("§6Finished restoring original blocks");
+            if (!cleanupUnplaceableBlocks.isEmpty()) {
+                SongPlayer.addChatMessage(String.format("§3%d §6blocks were not successfully restored"));
+            }
+        }
+    }
+    private void checkCleanupStatus() {
+        ClientWorld world = SongPlayer.MC.world;
+
+        cleanupPlaceList.clear();
+        cleanupBreakList.clear();
+        cleanupUnplaceableBlocks.clear();
+
+        for (BlockPos bp : originalBlocks.keySet()) {
+            BlockState actualBlockState = world.getBlockState(bp);
+            BlockState desiredBlockState = originalBlocks.get(bp);
+            if (actualBlockState != desiredBlockState) {
+                if (isPlaceable(desiredBlockState)) {
+                    cleanupPlaceList.add(bp);
+                }
+                if (!actualBlockState.isAir() && !actualBlockState.isLiquid()) {
+                    cleanupBreakList.add(bp);
+                }
+            }
+        }
+
+        cleanupBreakList = cleanupBreakList.stream()
+                .sorted((a, b) -> {
+                    // First sort by gravity
+                    boolean a_grav = SongPlayer.MC.world.getBlockState(a).getBlock() instanceof FallingBlock;
+                    boolean b_grav = SongPlayer.MC.world.getBlockState(b).getBlock() instanceof FallingBlock;
+                    if (a_grav && !b_grav) {
+                        return 1;
+                    } else if (!a_grav && b_grav) {
+                        return -1;
+                    }
+                    // Then sort by distance
+                    int a_dx = a.getX() - lastStage.position.getX();
+                    int a_dy = a.getY() - lastStage.position.getY();
+                    int a_dz = a.getZ() - lastStage.position.getZ();
+                    int b_dx = b.getX() - lastStage.position.getX();
+                    int b_dy = b.getY() - lastStage.position.getY();
+                    int b_dz = b.getZ() - lastStage.position.getZ();
+                    int a_dist = a_dx*a_dx + a_dy*a_dy + a_dz*a_dz;
+                    int b_dist = b_dx*b_dx + b_dy*b_dy + b_dz*b_dz;
+                    if (a_dist < b_dist) {
+                        return -1;
+                    } else if (a_dist > b_dist) {
+                        return 1;
+                    }
+                    // Finally sort by angle
+                    double a_angle = Math.atan2(a_dz, a_dx);
+                    double b_angle = Math.atan2(b_dz, b_dx);
+                    if (a_angle < b_angle) {
+                        return -1;
+                    } else if (a_angle > b_angle) {
+                        return 1;
+                    } else {
+                        return 0;
+                    }
+                })
+                .collect(Collectors.toCollection(LinkedList::new));
+
+        cleanupPlaceList = cleanupPlaceList.stream()
+                .sorted((a, b) -> {
+                    // First sort by gravity
+                    boolean a_grav = originalBlocks.get(a).getBlock() instanceof FallingBlock;
+                    boolean b_grav = originalBlocks.get(b).getBlock() instanceof FallingBlock;
+                    if (a_grav && !b_grav) {
+                        return -1;
+                    } else if (!a_grav && b_grav) {
+                        return 1;
+                    }
+                    // Then sort by distance
+                    int a_dx = a.getX() - lastStage.position.getX();
+                    int a_dy = a.getY() - lastStage.position.getY();
+                    int a_dz = a.getZ() - lastStage.position.getZ();
+                    int b_dx = b.getX() - lastStage.position.getX();
+                    int b_dy = b.getY() - lastStage.position.getY();
+                    int b_dz = b.getZ() - lastStage.position.getZ();
+                    int a_dist = a_dx*a_dx + a_dy*a_dy + a_dz*a_dz;
+                    int b_dist = b_dx*b_dx + b_dy*b_dy + b_dz*b_dz;
+                    if (a_dist < b_dist) {
+                        return -1;
+                    } else if (a_dist > b_dist) {
+                        return 1;
+                    }
+                    // Finally sort by angle
+                    double a_angle = Math.atan2(a_dz, a_dx);
+                    double b_angle = Math.atan2(b_dz, b_dx);
+                    if (a_angle < b_angle) {
+                        return 1;
+                    } else if (a_angle > b_angle) {
+                        return -1;
+                    } else {
+                        return 0;
+                    }
+                })
+                .collect(Collectors.toCollection(LinkedList::new));
+
+        cleanupPlaceList = cleanupPlaceList.reversed();
+        cleanupTotalBlocksToPlace = cleanupPlaceList.size();
+
+        for (BlockPos bp : cleanupPlaceList) {
+            System.out.println(world.getBlockState(bp).getBlock() + " " + originalBlocks.get(bp).getBlock());
+        }
+
+        boolean noNecessaryBreaks = cleanupBreakList.stream().allMatch(
+                bp -> world.getBlockState(bp).getBlock().getDefaultState().equals(originalBlocks.get(bp).getBlock().getDefaultState())
+        );
+        boolean noNecessaryPlacements = cleanupPlaceList.stream().allMatch(
+                bp -> bp.equals(lastStage.position)
+                || bp.equals(lastStage.position.up())
+                || world.getBlockState(bp).getBlock().getDefaultState().equals(originalBlocks.get(bp).getBlock().getDefaultState())
+        );
+        if (noNecessaryBreaks && noNecessaryPlacements) {
+            cleanupUnplaceableBlocks.addAll(cleanupPlaceList);
+            cleanupPlaceList.clear();
+        }
+    }
+    private void setCleanupProgressDisplay() {
+        MutableText buildText = Text.empty()
+                .append(Text.literal("Rebuilding original blocks | " ).formatted(Formatting.GOLD))
+                .append(Text.literal((cleanupTotalBlocksToPlace - cleanupPlaceList.size()) + "/" + cleanupTotalBlocksToPlace).formatted(Formatting.DARK_AQUA));
+        ProgressDisplay.getInstance().setText(buildText, Text.empty());
+    }
+
+    // Resets all internal states like currentSong, and songQueue, which stops all actions
+    public void reset() {
         currentSong = null;
         currentPlaylist = null;
         songQueue.clear();
         stage = null;
         buildSlot = -1;
         SongPlayer.removeFakePlayer();
+        cleaningUp = false;
+        dirty = false;
     }
-
+    public void restoreStateAndReset() {
-    public void restoreStateAndCleanUp() {
+        if (lastStage != null) {
-        if (stage != null) {
+            lastStage.movePlayerToStagePosition();
-            stage.movePlayerToStagePosition();
         }
         if (originalGamemode != SongPlayer.MC.interactionManager.getCurrentGameMode()) {
             if (originalGamemode == GameMode.CREATIVE) {
@@ -404,15 +619,38 @@ public class SongHandler {
             }
         }
         restoreBuildSlot();
-        cleanup();
+        reset();
+    }
+    public void partionResetAndCleanup() {
+        restoreBuildSlot();
+        currentSong = null;
+        currentPlaylist = null;
+        songQueue.clear();
+        stage = null;
+        buildSlot = -1;
+        startCleanup();
     }
 
+    // Runs every frame when player is not ingame
     public void onNotIngame() {
         currentSong = null;
         currentPlaylist = null;
         songQueue.clear();
     }
 
+    // Create stage if it doesn't exist and move the player to it
+    private void prepareStage() {
+        if (stage == null) {
+            stage = new Stage();
+            lastStage = stage;
+            originalBlocks.clear();
+            stage.movePlayerToStagePosition();
+        }
+        else {
+            stage.sendMovementPacketToStagePosition();
+        }
+    }
+
     private long lastCommandTime = System.currentTimeMillis();
     private String cachedCommand = null;
     private String cachedMessage = null;
@@ -468,10 +706,25 @@ public class SongHandler {
         inventory.main.set(slot, noteblockStack);
         SongPlayer.MC.interactionManager.clickCreativeStack(noteblockStack, 36 + slot);
     }
+    private void holdBlock(BlockState bs, int slot) {
+        PlayerInventory inventory = SongPlayer.MC.player.getInventory();
+        inventory.selectedSlot = slot;
+        ((ClientPlayerInteractionManagerAccessor) SongPlayer.MC.interactionManager).invokeSyncSelectedSlot();
+        ItemStack stack = new ItemStack(bs.getBlock());
+        Map<String, String> stateMap = new TreeMap<>();
+        for (Map.Entry<Property<?>, Comparable<?>> entry : bs.getEntries().entrySet()) {
+            Property<?> property = entry.getKey();
+            Comparable<?> value = entry.getValue();
+            stateMap.put(property.getName(), net.minecraft.util.Util.getValueAsString(property, value));
+        }
+        stack.set(DataComponentTypes.BLOCK_STATE, new BlockStateComponent(stateMap));
+        inventory.main.set(slot, stack);
+        SongPlayer.MC.interactionManager.clickCreativeStack(stack, 36 + slot);
+    }
     private void placeBlock(BlockPos bp) {
-        double fx = Math.max(0.0, Math.min(1.0, (stage.position.getX() + 0.5 - bp.getX())));
+        double fx = Math.max(0.0, Math.min(1.0, (lastStage.position.getX() + 0.5 - bp.getX())));
-        double fy = Math.max(0.0, Math.min(1.0, (stage.position.getY() + 0.0 - bp.getY())));
+        double fy = Math.max(0.0, Math.min(1.0, (lastStage.position.getY() + 0.0 - bp.getY())));
-        double fz = Math.max(0.0, Math.min(1.0, (stage.position.getZ() + 0.5 - bp.getZ())));
+        double fz = Math.max(0.0, Math.min(1.0, (lastStage.position.getZ() + 0.5 - bp.getZ())));
         fx += bp.getX();
         fy += bp.getY();
         fz += bp.getZ();
@@ -485,6 +738,33 @@ public class SongHandler {
     private void stopAttack() {
         SongPlayer.MC.interactionManager.cancelBlockBreaking();
     }
+    private void recordBlocks(Iterable<BlockPos> bpList) {
+        for (BlockPos bp : bpList) {
+            if (!originalBlocks.containsKey(bp)) {
+                BlockState bs = SongPlayer.MC.world.getBlockState(bp);
+                originalBlocks.put(bp, bs);
+            }
+        }
+    }
+    private void recordStageBlocks() {
+        recordBlocks(stage.requiredBreaks);
+        recordBlocks(stage.missingNotes
+                .stream()
+                .map(noteId -> stage.noteblockPositions.get(noteId))
+                .filter(Objects::nonNull)
+                .toList()
+        );
+    }
+    private boolean isPlaceable(BlockState bs) {
+        Block block = bs.getBlock();
+        if (bs.isAir() || bs.isLiquid()) {
+            return false;
+        } else if (block instanceof DoorBlock || block instanceof BedBlock) {
+            return false;
+        } else {
+            return true;
+        }
+    }
 
     private void doMovements(double lookX, double lookY, double lookZ) {
         if (Config.getConfig().swing) {
@@ -494,9 +774,9 @@ public class SongHandler {
             }
         }
         if (Config.getConfig().rotate) {
-            double d = lookX - (stage.position.getX() + 0.5);
+            double d = lookX - (lastStage.position.getX() + 0.5);
-            double e = lookY - (stage.position.getY() + SongPlayer.MC.player.getStandingEyeHeight());
+            double e = lookY - (lastStage.position.getY() + SongPlayer.MC.player.getStandingEyeHeight());
-            double f = lookZ - (stage.position.getZ() + 0.5);
+            double f = lookZ - (lastStage.position.getZ() + 0.5);
             double g = Math.sqrt(d * d + f * f);
             float pitch = MathHelper.wrapDegrees((float) (-(MathHelper.atan2(e, g) * 57.2957763671875)));
             float yaw = MathHelper.wrapDegrees((float) (MathHelper.atan2(f, d) * 57.2957763671875) - 90.0f);
@@ -506,7 +786,7 @@ public class SongHandler {
                 SongPlayer.fakePlayer.setHeadYaw(yaw);
             }
             SongPlayer.MC.player.networkHandler.getConnection().send(new PlayerMoveC2SPacket.Full(
-                    stage.position.getX() + 0.5, stage.position.getY(), stage.position.getZ() + 0.5,
+                    lastStage.position.getX() + 0.5, lastStage.position.getY(), lastStage.position.getZ() + 0.5,
                     yaw, pitch,
                     true));
         }
@@ -525,6 +805,6 @@ public class SongHandler {
     }
 
     public boolean isIdle() {
-        return currentSong == null && currentPlaylist == null && songQueue.isEmpty();
+        return currentSong == null && currentPlaylist == null && songQueue.isEmpty() && cleaningUp == false;
     }
 }

src/main/java/com/github/hhhzzzsss/songplayer/playing/Stage.java

@@ -468,4 +468,8 @@ public class Stage {
 		}
 		return false;
 	}
+
+	public Vec3d getOriginBottomCenter() {
+		return Vec3d.ofBottomCenter(position);
+	}
 }