From ae039cf14ec0d415fc67f1c9e67e7060b27123ae Mon Sep 17 00:00:00 2001
From: Adam <Jikoo.Games@gmail.com>
Date: Mon, 25 Feb 2019 04:30:59 -0500
Subject: [PATCH] Properly restore ArmorStands' gear (#64)

armor: Improve handling of armor stands, modernize
---
 .../CreeperHeal/block/CreeperArmorStand.java  | 63 +++++++++++--------
 1 file changed, 37 insertions(+), 26 deletions(-)

diff --git a/src/main/java/com/nitnelave/CreeperHeal/block/CreeperArmorStand.java b/src/main/java/com/nitnelave/CreeperHeal/block/CreeperArmorStand.java
index d4ebf4f..e8ad876 100644
--- a/src/main/java/com/nitnelave/CreeperHeal/block/CreeperArmorStand.java
+++ b/src/main/java/com/nitnelave/CreeperHeal/block/CreeperArmorStand.java
@@ -1,17 +1,19 @@
 package com.nitnelave.CreeperHeal.block;
 
+import com.nitnelave.CreeperHeal.CreeperHeal;
 import com.nitnelave.CreeperHeal.config.CreeperConfig;
 import com.nitnelave.CreeperHeal.utils.CreeperLog;
+import org.bukkit.Bukkit;
 import org.bukkit.Location;
 import org.bukkit.Material;
 import org.bukkit.World;
 import org.bukkit.block.Block;
 import org.bukkit.block.BlockFace;
 import org.bukkit.entity.ArmorStand;
+import org.bukkit.inventory.EntityEquipment;
 import org.bukkit.inventory.ItemStack;
 
-import java.util.ArrayList;
-import java.util.Collections;
+import java.util.Arrays;
 
 /**
  * Represents an ArmorStand
@@ -19,42 +21,57 @@
 public class CreeperArmorStand implements Replaceable
 {
     private final ArmorStand stand;
-    private ArrayList<ItemStack> contents = new ArrayList<ItemStack>();
+    private final ItemStack mainHand, offHand;
+    private final ItemStack[] contents;
     private boolean wasRemoved = false;
 
-    public CreeperArmorStand(ArmorStand stand)
+    CreeperArmorStand(ArmorStand stand)
     {
         this.stand = stand;
-        Collections.addAll(contents,
-                           stand.getBoots(), stand.getChestplate(), stand.getHelmet(), stand.getItemInHand(),
-                           stand.getLeggings());
-        CreeperLog.debug("Armor: " + contents);
+        EntityEquipment equipment = stand.getEquipment();
+        this.mainHand = equipment.getItemInMainHand();
+        this.offHand = equipment.getItemInOffHand();
+
+        this.contents = new ItemStack[equipment.getArmorContents().length];
+        System.arraycopy(equipment.getArmorContents(), 0, contents, 0, contents.length);
+        CreeperLog.debug("Armor: " + Arrays.toString(contents));
         remove();
     }
 
     @Override
     public boolean replace(boolean shouldDrop)
     {
-        ArmorStand s = getWorld()
-                .spawn(getLocation().getBlock().getRelative(BlockFace.DOWN).getLocation(), ArmorStand.class);
+        final ArmorStand s = getWorld().spawn(stand.getLocation(), ArmorStand.class);
         s.setArms(stand.hasArms());
         s.setBasePlate(stand.hasBasePlate());
         s.setBodyPose(stand.getBodyPose());
-        s.setBoots(stand.getBoots());
-        s.setChestplate(stand.getChestplate());
+        s.setCustomName(stand.getCustomName());
+        s.setCustomNameVisible(stand.isCustomNameVisible());
+        s.setGlowing(stand.isGlowing());
+        s.setGravity(stand.hasGravity());
         s.setHeadPose(stand.getHeadPose());
-        s.setHelmet(stand.getHelmet());
-        s.setItemInHand(stand.getItemInHand());
         s.setLeftArmPose(stand.getLeftArmPose());
         s.setRightArmPose(stand.getRightArmPose());
         s.setLeftLegPose(stand.getLeftLegPose());
         s.setRightLegPose(stand.getRightLegPose());
-        s.setLeggings(stand.getLeggings());
         s.setMarker(stand.isMarker());
         s.setSmall(stand.isSmall());
         s.setVisible(stand.isVisible());
 
-        s.teleport(stand.getLocation());
+        EntityEquipment equipment = s.getEquipment();
+        equipment.setArmorContents(contents);
+        equipment.setItemInMainHand(this.mainHand);
+        equipment.setItemInOffHand(this.offHand);
+
+        Bukkit.getScheduler().runTask(CreeperHeal.getInstance(), new Runnable()
+        {
+            @Override
+            public void run()
+            {
+                s.teleport(stand.getLocation());
+            }
+        });
+
         return true;
     }
 
@@ -100,9 +117,9 @@ public boolean drop(boolean forced)
         if (forced || CreeperConfig.shouldDrop())
         {
             getWorld().dropItemNaturally(getLocation(), new ItemStack(Material.ARMOR_STAND, 1));
-            for (ItemStack s : contents)
-                if (s.getType() != Material.AIR)
-                    getWorld().dropItemNaturally(getLocation(), s);
+            for (ItemStack itemStack : contents)
+                if (itemStack != null && itemStack.getType() != Material.AIR)
+                    getWorld().dropItemNaturally(getLocation(), itemStack);
             return true;
         }
         return false;
@@ -114,13 +131,7 @@ public void remove()
         if (!wasRemoved)
         {
             wasRemoved = true;
-            ItemStack air = new ItemStack(Material.AIR);
-            stand.setChestplate(air);
-            stand.setHelmet(air);
-            stand.setLeggings(air);
-            stand.setBoots(air);
-            stand.setItemInHand(air);
-
+            stand.getEquipment().clear();
             CreeperLog.debug("Removing armor, chestplate = " + stand.getChestplate().getType());
             stand.remove();
         }