From 4c589caef3e18bcf0c7875167bd48210ce12d469 Mon Sep 17 00:00:00 2001
From: Kli Kli <klikli@outlook.at>
Date: Thu, 1 Jun 2023 18:42:25 +0200
Subject: [PATCH] feat: allow divination rods to find blocks in same tag

and fix not finding any results if a block does not have a deepslate equivalent
---
 .../theurgy/item/DivinationRodItem.java       | 67 ++++++++++++++-----
 1 file changed, 49 insertions(+), 18 deletions(-)

diff --git a/src/main/java/com/klikli_dev/theurgy/item/DivinationRodItem.java b/src/main/java/com/klikli_dev/theurgy/item/DivinationRodItem.java
index 57c018b33..8c235bd1f 100644
--- a/src/main/java/com/klikli_dev/theurgy/item/DivinationRodItem.java
+++ b/src/main/java/com/klikli_dev/theurgy/item/DivinationRodItem.java
@@ -20,6 +20,7 @@
 import net.minecraft.client.renderer.item.ItemPropertyFunction;
 import net.minecraft.core.BlockPos;
 import net.minecraft.core.NonNullList;
+import net.minecraft.core.Registry;
 import net.minecraft.nbt.CompoundTag;
 import net.minecraft.network.chat.*;
 import net.minecraft.resources.ResourceLocation;
@@ -36,6 +37,7 @@
 import net.minecraft.world.item.context.UseOnContext;
 import net.minecraft.world.level.Level;
 import net.minecraft.world.level.block.Block;
+import net.minecraft.world.level.block.Blocks;
 import net.minecraft.world.level.block.state.BlockState;
 import net.minecraft.world.phys.Vec3;
 import net.minecraftforge.api.distmarker.Dist;
@@ -47,6 +49,7 @@
 
 import java.util.List;
 import java.util.Set;
+import java.util.stream.Collectors;
 
 public class DivinationRodItem extends Item {
 
@@ -73,6 +76,44 @@ public DivinationRodItem(Properties pProperties, Tier defaultTier, TagKey<Block>
         this.defaultAllowAttuning = defaultAllowAttuning;
     }
 
+    public static Set<Block> getScanTargetsForId(ResourceLocation linkedBlockId) {
+        //First: try to get a tag for the given block.
+        var tagKey = TagKey.create(Registry.BLOCK_REGISTRY, getOreTagFromBlockId(linkedBlockId));
+        var tag = ForgeRegistries.BLOCKS.tags().getTag(tagKey);
+
+        if (!tag.isEmpty())
+            return tag.stream().collect(Collectors.toSet());
+
+        //If no fitting tag succeeds, try to get block + deepslate variant
+        var block = ForgeRegistries.BLOCKS.getValue(linkedBlockId);
+        if (block != null) {
+            Block deepslateBlock = null;
+
+            //also search for deepslate ores
+            if (linkedBlockId.getPath().contains("_ore") && !linkedBlockId.getPath().contains("deepslate_")) {
+                var deepslateId = new ResourceLocation(linkedBlockId.getNamespace(), "deepslate_" + linkedBlockId.getPath());
+                deepslateBlock = ForgeRegistries.BLOCKS.getValue(deepslateId);
+            }
+
+            //finally, only add deepslate variant, if it is not air
+            return deepslateBlock != null && deepslateBlock != Blocks.AIR ? Set.of(block, deepslateBlock) : Set.of(block);
+        }
+
+        return Set.of();
+    }
+
+    public static ResourceLocation getOreTagFromBlockId(ResourceLocation blockId) {
+        var path = blockId.getPath();
+
+        String oreName = path
+                .replace("_ore", "")
+                .replace("ore_", "")
+                .replace("_deepslate", "")
+                .replace("deepslate_", "");
+
+        return new ResourceLocation("forge:ores/" + oreName);
+    }
+
     @Override
     public int getMaxDamage(ItemStack stack) {
         return stack.getOrCreateTag().getInt(TheurgyConstants.Nbt.Divination.SETTING_DURABILITY);
@@ -179,23 +220,13 @@ public InteractionResultHolder<ItemStack> use(Level level, Player player, Intera
 
                 if (level.isClientSide) {
                     var id = new ResourceLocation(stack.getTag().getString(TheurgyConstants.Nbt.Divination.LINKED_BLOCK_ID));
-                    var block = ForgeRegistries.BLOCKS.getValue(id);
-                    if (block != null) {
-                        Block deepslateBlock = null;
-
-                        //also search for deepslate ores
-                        if (id.getPath().contains("_ore") && !id.getPath().contains("deepslate_")) {
-                            var deepslateId = new ResourceLocation(id.getNamespace(), "deepslate_" + id.getPath());
-                            deepslateBlock = ForgeRegistries.BLOCKS.getValue(deepslateId);
-                        }
-
-                        ScanManager.get().beginScan(player,
-                                deepslateBlock != null ? Set.of(block, deepslateBlock) : Set.of(block),
-                                tag.getInt(TheurgyConstants.Nbt.Divination.SETTING_RANGE),
-                                tag.getInt(TheurgyConstants.Nbt.Divination.SETTING_DURATION)
-                        );
-                    }
 
+                    var blocks = getScanTargetsForId(id);
+                    ScanManager.get().beginScan(player,
+                            blocks,
+                            tag.getInt(TheurgyConstants.Nbt.Divination.SETTING_RANGE),
+                            tag.getInt(TheurgyConstants.Nbt.Divination.SETTING_DURATION)
+                    );
                 }
             } else if (!level.isClientSide) {
                 player.sendSystemMessage(Component.translatable(TheurgyConstants.I18n.Message.DIVINATION_ROD_NO_LINK));
@@ -210,7 +241,7 @@ public ItemStack finishUsingItem(ItemStack stack, Level level, LivingEntity enti
         if (!(entityLiving instanceof Player player))
             return stack;
 
-        if(stack.getDamageValue() >= stack.getMaxDamage()){
+        if (stack.getDamageValue() >= stack.getMaxDamage()) {
             //if in the last usage cycle the item was used up, we now actually break it to avoid over-use
             player.broadcastBreakEvent(player.getUsedItemHand());
             var item = stack.getItem();
@@ -446,7 +477,7 @@ public static class DistHelper {
         };
 
         public static void fillItemCategory(DivinationRodItem item, CreativeModeTab tab, NonNullList<ItemStack> items) {
-            if(tab != Theurgy.CREATIVE_MODE_TAB && tab != CreativeModeTab.TAB_SEARCH)
+            if (tab != Theurgy.CREATIVE_MODE_TAB && tab != CreativeModeTab.TAB_SEARCH)
                 return;
 
             var level = Minecraft.getInstance().level;