package net.minecraft.world.level.storage.loot.functions;

import com.google.gson.JsonDeserializationContext;
import net.minecraft.world.level.storage.loot.RandomIntGenerators;
import com.google.gson.JsonSerializationContext;
import com.google.gson.JsonObject;
import net.minecraft.resources.ResourceLocation;
import net.minecraft.world.level.storage.loot.LootContext;
import net.minecraft.world.item.ItemStack;
import net.minecraft.world.level.storage.loot.predicates.LootItemCondition;
import net.minecraft.world.level.storage.loot.RandomIntGenerator;

public class SetItemCountFunction extends LootItemConditionalFunction {
    private final RandomIntGenerator value;
    
    private SetItemCountFunction(final LootItemCondition[] arr, final RandomIntGenerator csd) {
        super(arr);
        this.value = csd;
    }
    
    public ItemStack run(final ItemStack ben, final LootContext crx) {
        ben.setCount(this.value.getInt(crx.getRandom()));
        return ben;
    }
    
    public static Builder<?> setCount(final RandomIntGenerator csd) {
        return LootItemConditionalFunction.simpleBuilder(arr -> new SetItemCountFunction(arr, csd));
    }
    
    public static class Serializer extends LootItemConditionalFunction.Serializer<SetItemCountFunction> {
        protected Serializer() {
            super(new ResourceLocation("set_count"), SetItemCountFunction.class);
        }
        
        @Override
        public void serialize(final JsonObject jsonObject, final SetItemCountFunction cto, final JsonSerializationContext jsonSerializationContext) {
            super.serialize(jsonObject, cto, jsonSerializationContext);
            jsonObject.add("count", RandomIntGenerators.serialize(cto.value, jsonSerializationContext));
        }
        
        @Override
        public SetItemCountFunction deserialize(final JsonObject jsonObject, final JsonDeserializationContext jsonDeserializationContext, final LootItemCondition[] arr) {
            final RandomIntGenerator csd5 = RandomIntGenerators.deserialize(jsonObject.get("count"), jsonDeserializationContext);
            return new SetItemCountFunction(arr, csd5, null);
        }
    }
}
