Add API to manipulate boss bar of entities and those created by commands

8 files changed, 227 insertions, 25 deletions

diff --git a/nms-patches/BossBattleCustom.patch b/nms-patches/BossBattleCustom.patch
new file mode 100644
index 00000000..cf8d10e9
--- /dev/null
+++ b/nms-patches/BossBattleCustom.patch
@@ -0,0 +1,30 @@
+--- a/net/minecraft/server/BossBattleCustom.java
++++ b/net/minecraft/server/BossBattleCustom.java
+@@ -8,12 +8,27 @@
+ import java.util.UUID;
+ import java.util.function.Consumer;
+ 
++// CraftBukkit start
++import org.bukkit.boss.KeyedBossBar;
++import org.bukkit.craftbukkit.boss.CraftKeyedBossbar;
++// CraftBukkit end
++
+ public class BossBattleCustom extends BossBattleServer {
+ 
+     private final MinecraftKey h;
+     private final Set<UUID> i = Sets.newHashSet();
+     private int j;
+     private int k = 100;
++    // CraftBukkit start
++    private KeyedBossBar bossBar;
++
++    public KeyedBossBar getBukkitEntity() {
++        if (bossBar == null) {
++            bossBar = new CraftKeyedBossbar(this);
++        }
++        return bossBar;
++    }
++    // CraftBukkit end
+ 
+     public BossBattleCustom(MinecraftKey minecraftkey, IChatBaseComponent ichatbasecomponent) {
+         super(ichatbasecomponent, BossBattle.BarColor.WHITE, BossBattle.BarStyle.PROGRESS);
diff --git a/nms-patches/EnderDragonBattle.patch b/nms-patches/EnderDragonBattle.patch
new file mode 100644
index 00000000..e5d6df48
--- /dev/null
+++ b/nms-patches/EnderDragonBattle.patch
@@ -0,0 +1,11 @@
+--- a/net/minecraft/server/EnderDragonBattle.java
++++ b/net/minecraft/server/EnderDragonBattle.java
+@@ -21,7 +21,7 @@
+ 
+     private static final Logger a = LogManager.getLogger();
+     private static final Predicate<Entity> b = IEntitySelector.a.and(IEntitySelector.a(0.0D, 128.0D, 0.0D, 192.0D));
+-    private final BossBattleServer c;
++    public final BossBattleServer c; // PAIL private -> public, rename bossBattleServer
+     private final WorldServer d;
+     private final List<Integer> e;
+     private final ShapeDetector f;
diff --git a/nms-patches/EntityWither.patch b/nms-patches/EntityWither.patch
index 77708f47..920308ca 100644
--- a/nms-patches/EntityWither.patch
+++ b/nms-patches/EntityWither.patch
@@ -13,6 +13,15 @@
  public class EntityWither extends EntityMonster implements IRangedEntity {
  
      private static final DataWatcherObject<Integer> a = DataWatcher.a(EntityWither.class, DataWatcherRegistry.b);
+@@ -19,7 +25,7 @@
+     private final int[] bI = new int[2];
+     private final int[] bJ = new int[2];
+     private int bK;
+-    private final BossBattleServer bL;
++    public final BossBattleServer bL; // PAIL private -> public, rename bossBattleServer
+     private static final Predicate<Entity> bM = (entity) -> {
+         return entity instanceof EntityLiving && ((EntityLiving) entity).getMonsterType() != EnumMonsterType.UNDEAD && ((EntityLiving) entity).df();
+     };
 @@ -181,13 +187,38 @@
          if (this.dz() > 0) {
              i = this.dz() - 1;
diff --git a/src/main/java/org/bukkit/craftbukkit/CraftServer.java b/src/main/java/org/bukkit/craftbukkit/CraftServer.java
index a1eacecf..e5be1ba6 100644
--- a/src/main/java/org/bukkit/craftbukkit/CraftServer.java
+++ b/src/main/java/org/bukkit/craftbukkit/CraftServer.java
@@ -44,6 +44,7 @@ import org.bukkit.boss.BarColor;
 import org.bukkit.boss.BarFlag;
 import org.bukkit.boss.BarStyle;
 import org.bukkit.boss.BossBar;
+import org.bukkit.boss.KeyedBossBar;
 import org.bukkit.command.Command;
 import org.bukkit.command.CommandException;
 import org.bukkit.command.CommandSender;
@@ -55,6 +56,7 @@ import org.bukkit.configuration.file.YamlConfiguration;
 import org.bukkit.configuration.serialization.ConfigurationSerialization;
 import org.bukkit.conversations.Conversable;
 import org.bukkit.craftbukkit.boss.CraftBossBar;
+import org.bukkit.craftbukkit.boss.CraftKeyedBossbar;
 import org.bukkit.craftbukkit.entity.CraftPlayer;
 import org.bukkit.craftbukkit.generator.CraftChunkData;
 import org.bukkit.craftbukkit.help.SimpleHelpMap;
@@ -153,6 +155,7 @@ import org.bukkit.craftbukkit.command.VanillaCommandWrapper;
 import org.bukkit.craftbukkit.inventory.util.CraftInventoryCreator;
 import org.bukkit.craftbukkit.tag.CraftBlockTag;
 import org.bukkit.craftbukkit.tag.CraftItemTag;
+import org.bukkit.craftbukkit.util.CraftChatMessage;
 import org.bukkit.craftbukkit.util.CraftNamespacedKey;
 import org.bukkit.event.server.ServerLoadEvent;
 import org.bukkit.event.server.TabCompleteEvent;
@@ -1733,6 +1736,53 @@ public final class CraftServer implements Server {
     }
 
     @Override
+    public KeyedBossBar createBossBar(NamespacedKey key, String title, BarColor barColor, BarStyle barStyle, BarFlag... barFlags) {
+        Preconditions.checkArgument(key != null, "key");
+
+        BossBattleCustom bossBattleCustom = getServer().aP().a(CraftNamespacedKey.toMinecraft(key), CraftChatMessage.fromString(title, true)[0]);
+        CraftKeyedBossbar craftKeyedBossbar = new CraftKeyedBossbar(bossBattleCustom);
+        craftKeyedBossbar.setColor(barColor);
+        craftKeyedBossbar.setStyle(barStyle);
+        for (BarFlag flag : barFlags) {
+            craftKeyedBossbar.addFlag(flag);
+        }
+
+        return craftKeyedBossbar;
+    }
+
+    @Override
+    public Iterator<KeyedBossBar> getBossBars() {
+        return Iterators.unmodifiableIterator(Iterators.transform(getServer().aP().b().iterator(), new Function<BossBattleCustom, org.bukkit.boss.KeyedBossBar>() { // PAIL: rename
+            @Override
+            public org.bukkit.boss.KeyedBossBar apply(BossBattleCustom bossBattleCustom) {
+                return bossBattleCustom.getBukkitEntity();
+            }
+        }));
+    }
+
+    @Override
+    public KeyedBossBar getBossBar(NamespacedKey key) {
+        Preconditions.checkArgument(key != null, "key");
+        net.minecraft.server.BossBattleCustom bossBattleCustom = getServer().aP().a(CraftNamespacedKey.toMinecraft(key));
+
+        return (bossBattleCustom == null) ? null : bossBattleCustom.getBukkitEntity();
+    }
+
+    @Override
+    public boolean removeBossBar(NamespacedKey key) {
+        Preconditions.checkArgument(key != null, "key");
+        net.minecraft.server.BossBattleCustomData bossBattleCustomData = getServer().aP();
+        net.minecraft.server.BossBattleCustom bossBattleCustom = bossBattleCustomData.a(CraftNamespacedKey.toMinecraft(key));
+
+        if (bossBattleCustom != null) {
+            bossBattleCustomData.a(bossBattleCustom);
+            return true;
+        }
+
+        return false;
+    }
+
+    @Override
     public Entity getEntity(UUID uuid) {
         Validate.notNull(uuid, "UUID cannot be null");
 
diff --git a/src/main/java/org/bukkit/craftbukkit/boss/CraftBossBar.java b/src/main/java/org/bukkit/craftbukkit/boss/CraftBossBar.java
index d9e54331..d39ee65c 100644
--- a/src/main/java/org/bukkit/craftbukkit/boss/CraftBossBar.java
+++ b/src/main/java/org/bukkit/craftbukkit/boss/CraftBossBar.java
@@ -14,29 +14,49 @@ import org.bukkit.craftbukkit.entity.CraftPlayer;
 import org.bukkit.craftbukkit.util.CraftChatMessage;
 import org.bukkit.entity.Player;
 
-import java.util.EnumSet;
+import java.util.HashMap;
 import java.util.List;
-import java.util.Set;
+import java.util.Map;
+import java.util.function.Consumer;
+import java.util.function.Supplier;
 
 public class CraftBossBar implements BossBar {
 
     private final BossBattleServer handle;
-    private final Set<BarFlag> flags;
-    private BarColor color;
-    private BarStyle style;
+    private Map<BarFlag, FlagContainer> flags;
 
     public CraftBossBar(String title, BarColor color, BarStyle style, BarFlag... flags) {
-        this.flags = flags.length > 0 ? EnumSet.of(flags[0], flags) : EnumSet.noneOf(BarFlag.class);
-        this.color = color;
-        this.style = style;
-
         handle = new BossBattleServer(
                 CraftChatMessage.fromString(title, true)[0],
                 convertColor(color),
                 convertStyle(style)
         );
 
-        updateFlags();
+        this.initialize();
+
+        for (BarFlag flag : flags) {
+            this.addFlag(flag);
+        }
+
+        this.setColor(color);
+        this.setStyle(style);
+    }
+
+    public CraftBossBar(BossBattleServer bossBattleServer) {
+        this.handle = bossBattleServer;
+        this.initialize();
+    }
+
+    private void initialize() {
+        this.flags = new HashMap<>();
+        this.flags.put(BarFlag.DARKEN_SKY, new FlagContainer(handle::n, handle::setDarkenSky));
+        this.flags.put(BarFlag.PLAY_BOSS_MUSIC, new FlagContainer(handle::o, handle::setPlayMusic));
+        this.flags.put(BarFlag.CREATE_FOG, new FlagContainer(handle::p, handle::setCreateFog));
+    }
+
+    private BarColor convertColor(BossBattle.BarColor color) {
+        BarColor bukkitColor = BarColor.valueOf(color.name());
+        return (bukkitColor == null) ? BarColor.WHITE : bukkitColor;
     }
 
     private BossBattle.BarColor convertColor(BarColor color) {
@@ -60,10 +80,20 @@ public class CraftBossBar implements BossBar {
         }
     }
 
-    private void updateFlags() {
-        handle.a(hasFlag(BarFlag.DARKEN_SKY));
-        handle.b(hasFlag(BarFlag.PLAY_BOSS_MUSIC));
-        handle.c(hasFlag(BarFlag.CREATE_FOG));
+    private BarStyle convertStyle(BossBattle.BarStyle style) {
+        switch (style) {
+            default:
+            case PROGRESS:
+                return BarStyle.SOLID;
+            case NOTCHED_6:
+                return BarStyle.SEGMENTED_6;
+            case NOTCHED_10:
+                return BarStyle.SEGMENTED_10;
+            case NOTCHED_12:
+                return BarStyle.SEGMENTED_12;
+            case NOTCHED_20:
+                return BarStyle.SEGMENTED_20;
+        }
     }
 
     @Override
@@ -79,48 +109,54 @@ public class CraftBossBar implements BossBar {
 
     @Override
     public BarColor getColor() {
-        return color;
+        return convertColor(handle.color);
     }
 
     @Override
     public void setColor(BarColor color) {
-        this.color = color;
         handle.color = convertColor(color);
         handle.sendUpdate(PacketPlayOutBoss.Action.UPDATE_STYLE);
     }
 
     @Override
     public BarStyle getStyle() {
-        return style;
+        return convertStyle(handle.style);
     }
 
     @Override
     public void setStyle(BarStyle style) {
-        this.style = style;
         handle.style = convertStyle(style);
         handle.sendUpdate(PacketPlayOutBoss.Action.UPDATE_STYLE);
     }
 
     @Override
     public void addFlag(BarFlag flag) {
-        flags.add(flag);
-        updateFlags();
+        FlagContainer flagContainer = flags.get(flag);
+        if (flagContainer != null) {
+            flagContainer.set.accept(true);
+        }
     }
 
     @Override
     public void removeFlag(BarFlag flag) {
-        flags.remove(flag);
-        updateFlags();
+        FlagContainer flagContainer = flags.get(flag);
+        if (flagContainer != null) {
+            flagContainer.set.accept(false);
+        }
     }
 
     @Override
     public boolean hasFlag(BarFlag flag) {
-        return flags.contains(flag);
+        FlagContainer flagContainer = flags.get(flag);
+        if (flagContainer != null) {
+            return flagContainer.get.get();
+        }
+        return false;
     }
 
     @Override
     public void setProgress(double progress) {
-    	Preconditions.checkArgument(progress >= 0.0 && progress <= 1.0, "Progress must be between 0.0 and 1.0 (%s)", progress);
+        Preconditions.checkArgument(progress >= 0.0 && progress <= 1.0, "Progress must be between 0.0 and 1.0 (%s)", progress);
         handle.setProgress((float) progress);
     }
 
@@ -174,4 +210,19 @@ public class CraftBossBar implements BossBar {
             removePlayer(player);
         }
     }
+
+    private class FlagContainer {
+
+        private Supplier<Boolean> get;
+        private Consumer<Boolean> set;
+
+        private FlagContainer(Supplier<Boolean> get, Consumer<Boolean> set) {
+            this.get = get;
+            this.set = set;
+        }
+    }
+
+    public BossBattleServer getHandle() {
+        return handle;
+    }
 }
diff --git a/src/main/java/org/bukkit/craftbukkit/boss/CraftKeyedBossbar.java b/src/main/java/org/bukkit/craftbukkit/boss/CraftKeyedBossbar.java
new file mode 100644
index 00000000..a35e6c20
--- /dev/null
+++ b/src/main/java/org/bukkit/craftbukkit/boss/CraftKeyedBossbar.java
@@ -0,0 +1,23 @@
+package org.bukkit.craftbukkit.boss;
+
+import net.minecraft.server.BossBattleCustom;
+import org.bukkit.NamespacedKey;
+import org.bukkit.boss.KeyedBossBar;
+import org.bukkit.craftbukkit.util.CraftNamespacedKey;
+
+public class CraftKeyedBossbar extends CraftBossBar implements KeyedBossBar {
+
+    public CraftKeyedBossbar(BossBattleCustom bossBattleCustom) {
+        super(bossBattleCustom);
+    }
+
+    @Override
+    public NamespacedKey getKey() {
+        return CraftNamespacedKey.fromMinecraft(getHandle().a());
+    }
+
+    @Override
+    public BossBattleCustom getHandle() {
+        return (BossBattleCustom) super.getHandle();
+    }
+}
diff --git a/src/main/java/org/bukkit/craftbukkit/entity/CraftEnderDragon.java b/src/main/java/org/bukkit/craftbukkit/entity/CraftEnderDragon.java
index d31a2069..97d5e628 100644
--- a/src/main/java/org/bukkit/craftbukkit/entity/CraftEnderDragon.java
+++ b/src/main/java/org/bukkit/craftbukkit/entity/CraftEnderDragon.java
@@ -9,14 +9,23 @@ import net.minecraft.server.DragonControllerPhase;
 import net.minecraft.server.EntityComplexPart;
 import net.minecraft.server.EntityEnderDragon;
 
+import org.bukkit.boss.BossBar;
 import org.bukkit.craftbukkit.CraftServer;
+import org.bukkit.craftbukkit.boss.CraftBossBar;
 import org.bukkit.entity.ComplexEntityPart;
 import org.bukkit.entity.EnderDragon;
 import org.bukkit.entity.EntityType;
 
 public class CraftEnderDragon extends CraftComplexLivingEntity implements EnderDragon {
+
+    private BossBar bossBar;
+
     public CraftEnderDragon(CraftServer server, EntityEnderDragon entity) {
         super(server, entity);
+
+        if (entity.ds() != null) {
+            this.bossBar = new CraftBossBar(entity.ds().c); // PAIL rename getEnderDragonBattle
+        }
     }
 
     public Set<ComplexEntityPart> getParts() {
@@ -60,4 +69,9 @@ public class CraftEnderDragon extends CraftComplexLivingEntity implements EnderD
     public static DragonControllerPhase getMinecraftPhase(Phase phase) {
         return DragonControllerPhase.getById(phase.ordinal());
     }
+
+    @Override
+    public BossBar getBossBar() {
+        return bossBar;
+    }
 }
diff --git a/src/main/java/org/bukkit/craftbukkit/entity/CraftWither.java b/src/main/java/org/bukkit/craftbukkit/entity/CraftWither.java
index fad3db8a..2b00b377 100644
--- a/src/main/java/org/bukkit/craftbukkit/entity/CraftWither.java
+++ b/src/main/java/org/bukkit/craftbukkit/entity/CraftWither.java
@@ -1,13 +1,22 @@
 package org.bukkit.craftbukkit.entity;
 
 import net.minecraft.server.EntityWither;
+import org.bukkit.boss.BossBar;
 import org.bukkit.craftbukkit.CraftServer;
-import org.bukkit.entity.Wither;
+import org.bukkit.craftbukkit.boss.CraftBossBar;
 import org.bukkit.entity.EntityType;
+import org.bukkit.entity.Wither;
 
 public class CraftWither extends CraftMonster implements Wither {
+
+    private BossBar bossBar;
+
     public CraftWither(CraftServer server, EntityWither entity) {
         super(server, entity);
+
+        if (entity.bL != null) {
+            this.bossBar = new CraftBossBar(entity.bL);
+        }
     }
 
     @Override
@@ -23,4 +32,9 @@ public class CraftWither extends CraftMonster implements Wither {
     public EntityType getType() {
         return EntityType.WITHER;
     }
+
+    @Override
+    public BossBar getBossBar() {
+        return bossBar;
+    }
 }