From 167f14bfe87dd3ee4c8d0f8e705342563e358a73 Mon Sep 17 00:00:00 2001
From: libraryaddict <libraryaddict115@yahoo.co.nz>
Date: Thu, 3 Apr 2014 02:02:58 +1300
Subject: [PATCH] Add support for disguiseNextEntity

---
 src/me/libraryaddict/disguise/DisguiseAPI.java   |  4 ++--
 .../disguise/utilities/DisguiseUtilities.java    | 16 +++++++++++++++-
 2 files changed, 17 insertions(+), 3 deletions(-)

diff --git a/src/me/libraryaddict/disguise/DisguiseAPI.java b/src/me/libraryaddict/disguise/DisguiseAPI.java
index 46f1dcf0..0ab7fe8c 100644
--- a/src/me/libraryaddict/disguise/DisguiseAPI.java
+++ b/src/me/libraryaddict/disguise/DisguiseAPI.java
@@ -78,7 +78,7 @@ public class DisguiseAPI {
             Field field = ReflectionManager.getNmsClass("Entity").getDeclaredField("entityCount");
             field.setAccessible(true);
             int id = field.getInt(null);
-            DisguiseUtilities.addDisguise(id, (TargetedDisguise) disguise);
+            DisguiseUtilities.addFutureDisguise(id, (TargetedDisguise) disguise);
         } catch (Exception ex) {
             ex.printStackTrace();
         }
@@ -128,7 +128,7 @@ public class DisguiseAPI {
     public static Disguise getDisguise(Player observer, Entity disguised) {
         if (disguised == null || observer == null)
             return null;
-        return DisguiseUtilities.getDisguise(observer, disguised.getUniqueId());
+        return DisguiseUtilities.getDisguise(observer, disguised);
     }
 
     /**
diff --git a/src/me/libraryaddict/disguise/utilities/DisguiseUtilities.java b/src/me/libraryaddict/disguise/utilities/DisguiseUtilities.java
index f684666d..9f3c9751 100644
--- a/src/me/libraryaddict/disguise/utilities/DisguiseUtilities.java
+++ b/src/me/libraryaddict/disguise/utilities/DisguiseUtilities.java
@@ -43,6 +43,14 @@ public class DisguiseUtilities {
     private static HashMap<UUID, Integer> selfDisguisesIds = new HashMap<UUID, Integer>();
     // Store the entity IDs instead of entitys because then I can disguise entitys even before they exist
     private static HashMap<UUID, HashSet<TargetedDisguise>> targetedDisguises = new HashMap<UUID, HashSet<TargetedDisguise>>();
+    private static HashMap<Integer, HashSet<TargetedDisguise>> futureDisguises = new HashMap<Integer, HashSet<TargetedDisguise>>();
+
+    public static void addFutureDisguise(int entityId, TargetedDisguise disguise) {
+        if (!futureDisguises.containsKey(entityId)) {
+            futureDisguises.put(entityId, new HashSet<TargetedDisguise>());
+        }
+        futureDisguises.get(entityId).add(disguise);
+    }
 
     public static void addDisguise(UUID entityId, TargetedDisguise disguise) {
         if (!getDisguises().containsKey(entityId)) {
@@ -163,7 +171,13 @@ public class DisguiseUtilities {
         }
     }
 
-    public static TargetedDisguise getDisguise(Player observer, UUID entityId) {
+    public static TargetedDisguise getDisguise(Player observer, Entity entity) {
+        UUID entityId = entity.getUniqueId();
+        if (futureDisguises.containsKey(entity.getEntityId())) {
+            for (TargetedDisguise disguise : futureDisguises.remove(entity.getEntityId())) {
+                addDisguise(entityId, disguise);
+            }
+        }
         if (getDisguises().containsKey(entityId)) {
             for (TargetedDisguise disguise : getDisguises().get(entityId)) {
                 if (disguise.canSee(observer)) {