From 964671317ad2d32f27e6993fcdf0a7f283a40d31 Mon Sep 17 00:00:00 2001
From: Frank <developer@ambertation.de>
Date: Wed, 20 Dec 2023 16:20:08 +0100
Subject: [PATCH] [Change] Add Flight Duration to Rockets from Village Loot
 (#309)

---
 .../loot_tables/chests/end_village_bonus_loot.json     |  4 ++++
 .../betterend/recipes/EndChestLootTableProvider.java   | 10 +++++++++-
 2 files changed, 13 insertions(+), 1 deletion(-)

diff --git a/src/main/generated/data/betterend/loot_tables/chests/end_village_bonus_loot.json b/src/main/generated/data/betterend/loot_tables/chests/end_village_bonus_loot.json
index db766a74..fc4f04f6 100644
--- a/src/main/generated/data/betterend/loot_tables/chests/end_village_bonus_loot.json
+++ b/src/main/generated/data/betterend/loot_tables/chests/end_village_bonus_loot.json
@@ -138,6 +138,10 @@
         {
           "type": "minecraft:item",
           "functions": [
+            {
+              "function": "minecraft:set_nbt",
+              "tag": "{Fireworks:{Flight:3b}}"
+            },
             {
               "add": false,
               "count": {
diff --git a/src/main/java/org/betterx/datagen/betterend/recipes/EndChestLootTableProvider.java b/src/main/java/org/betterx/datagen/betterend/recipes/EndChestLootTableProvider.java
index d4fef8bb..58cabd5f 100644
--- a/src/main/java/org/betterx/datagen/betterend/recipes/EndChestLootTableProvider.java
+++ b/src/main/java/org/betterx/datagen/betterend/recipes/EndChestLootTableProvider.java
@@ -6,6 +6,7 @@ import org.betterx.betterend.registry.EndItems;
 import org.betterx.betterend.registry.EndTemplates;
 import org.betterx.betterend.util.LootTableUtil;
 
+import net.minecraft.nbt.CompoundTag;
 import net.minecraft.resources.ResourceLocation;
 import net.minecraft.world.item.Items;
 import net.minecraft.world.level.block.Blocks;
@@ -15,6 +16,7 @@ import net.minecraft.world.level.storage.loot.entries.LootItem;
 import net.minecraft.world.level.storage.loot.functions.EnchantWithLevelsFunction;
 import net.minecraft.world.level.storage.loot.functions.SetItemCountFunction;
 import net.minecraft.world.level.storage.loot.functions.SetItemDamageFunction;
+import net.minecraft.world.level.storage.loot.functions.SetNbtFunction;
 import net.minecraft.world.level.storage.loot.parameters.LootContextParamSets;
 import net.minecraft.world.level.storage.loot.predicates.LootItemRandomChanceCondition;
 import net.minecraft.world.level.storage.loot.providers.number.ConstantValue;
@@ -26,11 +28,15 @@ import net.fabricmc.fabric.api.datagen.v1.provider.SimpleFabricLootTableProvider
 import java.util.function.BiConsumer;
 
 public class EndChestLootTableProvider extends SimpleFabricLootTableProvider {
+    private final CompoundTag FIREWORK_ROCKET_TAG = new CompoundTag();
 
     public EndChestLootTableProvider(
             FabricDataOutput output
     ) {
         super(output, LootContextParamSets.CHEST);
+        CompoundTag FIREWORKS_TAG = new CompoundTag();
+        FIREWORK_ROCKET_TAG.put("Fireworks", FIREWORKS_TAG);
+        FIREWORKS_TAG.putByte("Flight", (byte) 3);
     }
 
 
@@ -316,7 +322,9 @@ public class EndChestLootTableProvider extends SimpleFabricLootTableProvider {
                 ).add(LootItem
                         .lootTableItem(Items.FIREWORK_ROCKET)
                         .setWeight(8)
-                        .apply(SetItemCountFunction.setCount(UniformGenerator.between(4, 8)))
+                        .apply(SetNbtFunction.setTag(FIREWORK_ROCKET_TAG))
+                        .apply(SetItemCountFunction
+                                .setCount(UniformGenerator.between(4, 8)))
                 )
                 .when(LootItemRandomChanceCondition.randomChance(chance));
     }