package info.u_team.u_team_core.intern.loot;

import com.google.gson.JsonDeserializationContext;
import com.google.gson.JsonObject;

import info.u_team.u_team_core.blockentity.UBlockEntity;
import info.u_team.u_team_core.intern.init.UCoreLootFunctions;
import net.minecraft.nbt.CompoundTag;
import net.minecraft.world.item.ItemStack;
import net.minecraft.world.level.block.entity.BlockEntity;
import net.minecraft.world.level.storage.loot.LootContext;
import net.minecraft.world.level.storage.loot.functions.LootItemConditionalFunction;
import net.minecraft.world.level.storage.loot.functions.LootItemFunctionType;
import net.minecraft.world.level.storage.loot.parameters.LootContextParams;
import net.minecraft.world.level.storage.loot.predicates.LootItemCondition;

public class SetBlockEntityNBTLootFunction extends LootItemConditionalFunction {
	
	private SetBlockEntityNBTLootFunction(LootItemCondition[] conditions) {
		super(conditions);
	}
	
	@Override
	public ItemStack run(ItemStack stack, LootContext context) {
		if (context.hasParam(LootContextParams.BLOCK_ENTITY)) {
			final BlockEntity blockEntity = context.getParamOrNull(LootContextParams.BLOCK_ENTITY);
			final CompoundTag compound;
			if (blockEntity instanceof final UBlockEntity uBlockEntity) {
				compound = new CompoundTag();
				uBlockEntity.saveNBT(compound);
			} else {
				compound = blockEntity.saveWithId();
			}
			if (!compound.isEmpty()) {
				stack.addTagElement("BlockEntityTag", compound);
			}
		}
		return stack;
	}
	
	public static LootItemConditionalFunction.Builder<?> builder() {
		return simpleBuilder((conditions) -> new SetBlockEntityNBTLootFunction(conditions));
	}
	
	@Override
	public LootItemFunctionType getType() {
		return UCoreLootFunctions.SET_BLOCKENTITY_NBT.get();
	}
	
	public static class Serializer extends LootItemConditionalFunction.Serializer<SetBlockEntityNBTLootFunction> {
		
		@Override
		public SetBlockEntityNBTLootFunction deserialize(JsonObject object, JsonDeserializationContext deserializationContext, LootItemCondition[] conditions) {
			return new SetBlockEntityNBTLootFunction(conditions);
		}
	}
}

@@ -161,133 +161,38 @@
 ore.
-blockentity.UBlockEntity;%0Aimport info.u_team.u_team_core.intern.init.UCoreLootFunctions;%0Aimport net.minecraft.nbt.CompoundTag
+intern.init.UCoreLootFunctions
 ;%0Aim
@@ -749,16 +749,147 @@
 ion %7B%0A%09%0A
+%09public static LootItemConditionalFunction.Builder%3C?%3E builder() %7B%0A%09%09return simpleBuilder(SetBlockEntityNBTLootFunction::new);%0A%09%7D%0A%09%0A
 %09private
@@ -1208,493 +1208,57 @@
 %0A%09%09%09
-final CompoundTag compound;%0A%09%09%09if (blockEntity instanceof final UBlockEntity uBlockEntity) %7B%0A%09%09%09%09compound = new CompoundTag();%0A%09%09%09%09uBlockEntity.saveNBT(compound);%0A%09%09%09%7D else %7B%0A%09%09%09%09compound = blockEntity.saveWithId();%0A%09%09%09%7D%0A%09%09%09if (!compound.isEmpty()) %7B%0A%09%09%09%09stack.addTagElement(%22BlockEntityTag%22, compound);%0A%09%09%09%7D%0A%09%09%7D%0A%09%09return stack;%0A%09%7D%0A%09%0A%09public static LootItemConditionalFunction.Builder%3C?%3E builder() %7B%0A%09%09return simpleBuilder((conditions) -%3E new SetBlockEntityNBTLootFunction(conditions))
+blockEntity.saveToItem(stack);%0A%09%09%7D%0A%09%09return stack
 ;%0A%09%7D

Update block entity nbt function to be easer accessible