From fa4c577f7e90f0d5b3b26a047052955c05fde461 Mon Sep 17 00:00:00 2001
From: libraryaddict <github@lib.co.nz>
Date: Wed, 22 Apr 2020 14:20:52 +1200
Subject: [PATCH] More disguiseNextEntity fixing

---
 .../disguise/utilities/DisguiseUtilities.java | 64 +++++++++++--------
 .../packetlisteners/PacketListenerMain.java   | 17 +----
 .../packetlisteners/PacketListenerSounds.java |  2 +-
 3 files changed, 40 insertions(+), 43 deletions(-)

diff --git a/src/main/java/me/libraryaddict/disguise/utilities/DisguiseUtilities.java b/src/main/java/me/libraryaddict/disguise/utilities/DisguiseUtilities.java
index 74a7e7c6..3bd48811 100644
--- a/src/main/java/me/libraryaddict/disguise/utilities/DisguiseUtilities.java
+++ b/src/main/java/me/libraryaddict/disguise/utilities/DisguiseUtilities.java
@@ -428,37 +428,49 @@ public class DisguiseUtilities {
         }
     }
 
-    public static void addFutureDisguise(final int entityId, final TargetedDisguise disguise) {
-        if (!futureDisguises.containsKey(entityId)) {
-            futureDisguises.put(entityId, new HashSet<TargetedDisguise>());
+    public static void onFutureDisguise(Entity entity) {
+        if (!getFutureDisguises().containsKey(entity.getEntityId())) {
+            return;
         }
 
-        futureDisguises.get(entityId).add(disguise);
+        UUID uniqueId = entity.getUniqueId();
+
+        for (TargetedDisguise disguise : getFutureDisguises().remove(entity.getEntityId())) {
+            addDisguise(uniqueId, disguise);
+        }
+    }
+
+    public static void addFutureDisguise(final int entityId, final TargetedDisguise disguise) {
+        if (!getFutureDisguises().containsKey(entityId)) {
+            getFutureDisguises().put(entityId, new HashSet<>());
+        }
+
+        getFutureDisguises().get(entityId).add(disguise);
 
         final BukkitRunnable runnable = new BukkitRunnable() {
             @Override
             public void run() {
-                if (futureDisguises.containsKey(entityId) && futureDisguises.get(entityId).contains(disguise)) {
-                    for (World world : Bukkit.getWorlds()) {
-                        for (Entity entity : world.getEntities()) {
-                            if (entity.getEntityId() == entityId) {
-                                UUID uniqueId = entity.getUniqueId();
+                if (!getFutureDisguises().containsKey(entityId) ||
+                        !getFutureDisguises().get(entityId).contains(disguise)) {
+                    return;
+                }
 
-                                for (TargetedDisguise disguise : futureDisguises.remove(entityId)) {
-                                    addDisguise(uniqueId, disguise);
-                                }
-
-                                return;
-                            }
+                for (World world : Bukkit.getWorlds()) {
+                    for (Entity entity : world.getEntities()) {
+                        if (entity.getEntityId() != entityId) {
+                            continue;
                         }
-                    }
 
-                    futureDisguises.get(entityId).remove(disguise);
-
-                    if (futureDisguises.get(entityId).isEmpty()) {
-                        futureDisguises.remove(entityId);
+                        onFutureDisguise(entity);
+                        return;
                     }
                 }
+
+                getFutureDisguises().get(entityId).remove(disguise);
+
+                if (getFutureDisguises().get(entityId).isEmpty()) {
+                    getFutureDisguises().remove(entityId);
+                }
             }
         };
 
@@ -2227,14 +2239,14 @@ public class DisguiseUtilities {
         }
 
         if (getFutureDisguises().containsKey(entityId)) {
-            HashSet<TargetedDisguise> hashSet = getFutureDisguises().get(entityId);
+            for (World world : Bukkit.getWorlds()) {
+                for (Entity entity : world.getEntities()) {
+                    if (entity.getEntityId() != entityId) {
+                        continue;
+                    }
 
-            for (TargetedDisguise dis : hashSet) {
-                if (!dis.canSee(observer) || !dis.isDisguiseInUse()) {
-                    continue;
+                    onFutureDisguise(entity);
                 }
-
-                return dis;
             }
         }
 
diff --git a/src/main/java/me/libraryaddict/disguise/utilities/packets/packetlisteners/PacketListenerMain.java b/src/main/java/me/libraryaddict/disguise/utilities/packets/packetlisteners/PacketListenerMain.java
index 82a029ef..371986ac 100644
--- a/src/main/java/me/libraryaddict/disguise/utilities/packets/packetlisteners/PacketListenerMain.java
+++ b/src/main/java/me/libraryaddict/disguise/utilities/packets/packetlisteners/PacketListenerMain.java
@@ -39,25 +39,10 @@ public class PacketListenerMain extends PacketAdapter {
 
         final Disguise disguise = DisguiseUtilities.getDisguise(observer, entityId);
 
-        if (disguise == null) {
-            return;
-        }
-
-        if (disguise.getEntity() == null) {
-            for (Entity e : observer.getWorld().getEntities()) {
-                if (e.getEntityId() != entityId) {
-                    continue;
-                }
-
-                disguise.setEntity(e);
-                break;
-            }
-        }
-
         // If the entity is the same as the sender. Don't disguise!
         // Prevents problems and there is no advantage to be gained.
         // Or if they are null and there's no disguise
-        if (disguise.getEntity() == observer) {
+        if (disguise == null || disguise.getEntity() == observer) {
             return;
         }
 
diff --git a/src/main/java/me/libraryaddict/disguise/utilities/packets/packetlisteners/PacketListenerSounds.java b/src/main/java/me/libraryaddict/disguise/utilities/packets/packetlisteners/PacketListenerSounds.java
index 166f8cf6..593840bd 100644
--- a/src/main/java/me/libraryaddict/disguise/utilities/packets/packetlisteners/PacketListenerSounds.java
+++ b/src/main/java/me/libraryaddict/disguise/utilities/packets/packetlisteners/PacketListenerSounds.java
@@ -257,7 +257,7 @@ public class PacketListenerSounds extends PacketAdapter {
 
             Entity entity = disguise.getEntity();
 
-            if (disguise != null && !disguise.getType().isPlayer() &&
+            if (!disguise.getType().isPlayer() &&
                     (disguise.isSelfDisguiseSoundsReplaced() || entity != event.getPlayer())) {
                 DisguiseSound disSound = DisguiseSound.getType(entity.getType().name());