Implemented enchantments and enchantmentLevels conditions

2 files changed, 121 insertions, 0 deletions

diff --git a/defaults/src/main/java/shcm/shsupercm/fabric/citresewn/defaults/cit/conditions/ConditionEnchantmentLevels.java b/defaults/src/main/java/shcm/shsupercm/fabric/citresewn/defaults/cit/conditions/ConditionEnchantmentLevels.java
new file mode 100644
index 0000000..da6b4f4
--- /dev/null
+++ b/defaults/src/main/java/shcm/shsupercm/fabric/citresewn/defaults/cit/conditions/ConditionEnchantmentLevels.java
@@ -0,0 +1,62 @@
+package shcm.shsupercm.fabric.citresewn.defaults.cit.conditions;
+
+import io.shcm.shsupercm.fabric.fletchingtable.api.Entrypoint;
+import net.minecraft.util.Identifier;
+import shcm.shsupercm.fabric.citresewn.api.CITConditionContainer;
+import shcm.shsupercm.fabric.citresewn.cit.CITCondition;
+import shcm.shsupercm.fabric.citresewn.cit.CITContext;
+import shcm.shsupercm.fabric.citresewn.cit.builtin.conditions.IntegerCondition;
+import shcm.shsupercm.fabric.citresewn.cit.builtin.conditions.ListCondition;
+
+import java.util.Arrays;
+import java.util.HashSet;
+import java.util.Map;
+import java.util.Set;
+
+public class ConditionEnchantmentLevels extends ListCondition<ConditionEnchantmentLevels.EnchantmentLevelCondition> {
+    @Entrypoint(CITConditionContainer.ENTRYPOINT)
+    public static final CITConditionContainer<ConditionEnchantmentLevels> CONTAINER = new CITConditionContainer<>(ConditionEnchantmentLevels.class, ConditionEnchantmentLevels::new,
+            "enchantment_levels", "enchantmentLevels");
+
+    protected Set<Identifier> enchantments = null;
+
+    public ConditionEnchantmentLevels() {
+        super(EnchantmentLevelCondition.class, EnchantmentLevelCondition::new);
+    }
+
+    @Override
+    public Set<Class<? extends CITCondition>> siblingConditions() {
+        return Set.of(ConditionEnchantments.class);
+    }
+
+    @Override
+    public <T extends CITCondition> T modifySibling(T sibling) {
+        if (sibling instanceof ConditionEnchantments conditionEnchantments) {
+            if (enchantments == null) {
+                enchantments = new HashSet<>();
+                for (EnchantmentLevelCondition subCondition : this.conditions)
+                    subCondition.enchantments = enchantments;
+            }
+            enchantments.addAll(Arrays.asList(conditionEnchantments.getEnchantments()));
+        }
+
+        return sibling;
+    }
+
+    protected static class EnchantmentLevelCondition extends IntegerCondition {
+        protected Set<Identifier> enchantments = null;
+
+        protected EnchantmentLevelCondition() {
+            super(true, false, false);
+        }
+
+        @Override
+        public boolean test(CITContext context) {
+            for (Map.Entry<Identifier, Integer> entry : context.enchantments().entrySet())
+                if ((enchantments == null || enchantments.contains(entry.getKey())) && entry.getValue() != null && (range ? min <= entry.getValue() && entry.getValue() <= max : entry.getValue() == min))
+                    return true;
+
+            return false;
+        }
+    }
+}
diff --git a/defaults/src/main/java/shcm/shsupercm/fabric/citresewn/defaults/cit/conditions/ConditionEnchantments.java b/defaults/src/main/java/shcm/shsupercm/fabric/citresewn/defaults/cit/conditions/ConditionEnchantments.java
new file mode 100644
index 0000000..301f9c1
--- /dev/null
+++ b/defaults/src/main/java/shcm/shsupercm/fabric/citresewn/defaults/cit/conditions/ConditionEnchantments.java
@@ -0,0 +1,59 @@
+package shcm.shsupercm.fabric.citresewn.defaults.cit.conditions;
+
+import io.shcm.shsupercm.fabric.fletchingtable.api.Entrypoint;
+import net.minecraft.util.Identifier;
+import net.minecraft.util.registry.Registry;
+import shcm.shsupercm.fabric.citresewn.api.CITConditionContainer;
+import shcm.shsupercm.fabric.citresewn.cit.CITCondition;
+import shcm.shsupercm.fabric.citresewn.cit.CITContext;
+import shcm.shsupercm.fabric.citresewn.cit.builtin.conditions.IdentifierCondition;
+import shcm.shsupercm.fabric.citresewn.cit.builtin.conditions.ListCondition;
+import shcm.shsupercm.fabric.citresewn.ex.CITParsingException;
+import shcm.shsupercm.fabric.citresewn.pack.format.PropertyGroup;
+import shcm.shsupercm.fabric.citresewn.pack.format.PropertyValue;
+
+import java.util.Set;
+
+public class ConditionEnchantments extends ListCondition<ConditionEnchantments.EnchantmentCondition> {
+    @Entrypoint(CITConditionContainer.ENTRYPOINT)
+    public static final CITConditionContainer<ConditionEnchantments> CONTAINER = new CITConditionContainer<>(ConditionEnchantments.class, ConditionEnchantments::new,
+            "enchantments", "enchantmentIDs");
+
+    public ConditionEnchantments() {
+        super(EnchantmentCondition.class, EnchantmentCondition::new);
+    }
+
+    public Identifier[] getEnchantments() {
+        Identifier[] enchantments = new Identifier[this.conditions.length];
+
+        for (int i = 0; i < this.conditions.length; i++)
+            enchantments[i] = this.conditions[i].getValue(null);
+
+        return enchantments;
+    }
+
+    @Override
+    public Set<Class<? extends CITCondition>> siblingConditions() {
+        return Set.of(ConditionEnchantmentLevels.class);
+    }
+
+    protected static class EnchantmentCondition extends IdentifierCondition {
+        @Override
+        public void load(PropertyValue value, PropertyGroup properties) throws CITParsingException {
+            super.load(value, properties);
+
+            if (!Registry.ENCHANTMENT.containsId(this.value))
+                warn(this.value + " is not in the enchantment registry", value, properties);
+        }
+
+        @Override
+        public boolean test(CITContext context) {
+            return context.enchantments().containsKey(this.value);
+        }
+
+        @Override
+        protected Identifier getValue(CITContext context) {
+            return this.value;
+        }
+    }
+}