From 99dff66345517b26f259019b3665cabb4744e9a4 Mon Sep 17 00:00:00 2001
From: simibubi <31564874+simibubi@users.noreply.github.com>
Date: Sun, 22 Oct 2023 12:17:14 +0200
Subject: [PATCH] Netherite Diving PR Tweaks

---
 .../equipment/armor/NetheriteDivingHandler.java    | 14 ++++++++++++--
 .../create/foundation/events/ClientEvents.java     |  3 +--
 2 files changed, 13 insertions(+), 4 deletions(-)

diff --git a/src/main/java/com/simibubi/create/content/equipment/armor/NetheriteDivingHandler.java b/src/main/java/com/simibubi/create/content/equipment/armor/NetheriteDivingHandler.java
index dda758712..41b09e7f6 100644
--- a/src/main/java/com/simibubi/create/content/equipment/armor/NetheriteDivingHandler.java
+++ b/src/main/java/com/simibubi/create/content/equipment/armor/NetheriteDivingHandler.java
@@ -1,5 +1,7 @@
 package com.simibubi.create.content.equipment.armor;
 
+import com.simibubi.create.AllTags.AllItemTags;
+
 import net.minecraft.nbt.CompoundTag;
 import net.minecraft.world.entity.EquipmentSlot;
 import net.minecraft.world.entity.LivingEntity;
@@ -26,13 +28,13 @@ public final class NetheriteDivingHandler {
 		ItemStack to = event.getTo();
 
 		if (slot == EquipmentSlot.HEAD) {
-			if (to.getItem() instanceof DivingHelmetItem && isNetheriteArmor(to)) {
+			if (isNetheriteDivingHelmet(to)) {
 				setBit(entity, slot);
 			} else {
 				clearBit(entity, slot);
 			}
 		} else if (slot == EquipmentSlot.CHEST) {
-			if (to.getItem() instanceof BacktankItem && isNetheriteArmor(to) && BacktankUtil.hasAirRemaining(to)) {
+			if (isNetheriteBacktank(to) && BacktankUtil.hasAirRemaining(to)) {
 				setBit(entity, slot);
 			} else {
 				clearBit(entity, slot);
@@ -46,6 +48,14 @@ public final class NetheriteDivingHandler {
 		}
 	}
 
+	public static boolean isNetheriteDivingHelmet(ItemStack stack) {
+		return stack.getItem() instanceof DivingHelmetItem && isNetheriteArmor(stack);
+	}
+	
+	public static boolean isNetheriteBacktank(ItemStack stack) {
+		return stack.is(AllItemTags.PRESSURIZED_AIR_SOURCES.tag) && isNetheriteArmor(stack);
+	}
+	
 	public static boolean isNetheriteArmor(ItemStack stack) {
 		return stack.getItem() instanceof ArmorItem armorItem && armorItem.getMaterial() == ArmorMaterials.NETHERITE;
 	}
diff --git a/src/main/java/com/simibubi/create/foundation/events/ClientEvents.java b/src/main/java/com/simibubi/create/foundation/events/ClientEvents.java
index 7510c512e..c373e6d9b 100644
--- a/src/main/java/com/simibubi/create/foundation/events/ClientEvents.java
+++ b/src/main/java/com/simibubi/create/foundation/events/ClientEvents.java
@@ -3,7 +3,6 @@ package com.simibubi.create.foundation.events;
 import com.mojang.blaze3d.systems.RenderSystem;
 import com.mojang.blaze3d.vertex.PoseStack;
 import com.simibubi.create.AllFluids;
-import com.simibubi.create.AllItems;
 import com.simibubi.create.AllPackets;
 import com.simibubi.create.Create;
 import com.simibubi.create.CreateClient;
@@ -317,7 +316,7 @@ public class ClientEvents {
 				event.scaleFarPlaneDistance(6.25f);
 				event.setCanceled(true);
 				return;
-			} else if (FluidHelper.isLava(fluid) && NetheriteDivingHandler.isNetheriteArmor(divingHelmet)) {
+			} else if (FluidHelper.isLava(fluid) && NetheriteDivingHandler.isNetheriteDivingHelmet(divingHelmet)) {
 				event.setNearPlaneDistance(-4.0f);
 				event.setFarPlaneDistance(20.0f);
 				event.setCanceled(true);