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

import net.minecraft.util.GsonHelper;
import com.google.gson.JsonDeserializationContext;
import com.google.gson.JsonSerializationContext;
import com.google.gson.JsonObject;
import net.minecraft.resources.ResourceLocation;
import net.minecraft.world.item.enchantment.EnchantmentHelper;
import net.minecraft.world.entity.LivingEntity;
import net.minecraft.world.level.storage.loot.LootContext;
import net.minecraft.world.item.ItemStack;
import net.minecraft.world.entity.Entity;
import com.google.common.collect.ImmutableSet;
import net.minecraft.world.level.storage.loot.parameters.LootContextParams;
import net.minecraft.world.level.storage.loot.parameters.LootContextParam;
import java.util.Set;
import net.minecraft.world.level.storage.loot.predicates.LootItemCondition;
import net.minecraft.world.level.storage.loot.RandomValueBounds;

public class LootingEnchantFunction extends LootItemConditionalFunction {
    private final RandomValueBounds value;
    private final int limit;
    
    private LootingEnchantFunction(final LootItemCondition[] arr, final RandomValueBounds csf, final int integer) {
        super(arr);
        this.value = csf;
        this.limit = integer;
    }
    
    @Override
    public Set<LootContextParam<?>> getReferencedContextParams() {
        return ImmutableSet.of(LootContextParams.KILLER_ENTITY);
    }
    
    private boolean hasLimit() {
        return this.limit > 0;
    }
    
    public ItemStack run(final ItemStack ben, final LootContext crx) {
        final Entity akq4 = crx.<Entity>getParamOrNull(LootContextParams.KILLER_ENTITY);
        if (akq4 instanceof LivingEntity) {
            final int integer5 = EnchantmentHelper.getMobLooting((LivingEntity)akq4);
            if (integer5 == 0) {
                return ben;
            }
            final float float6 = integer5 * this.value.getFloat(crx.getRandom());
            ben.grow(Math.round(float6));
            if (this.hasLimit() && ben.getCount() > this.limit) {
                ben.setCount(this.limit);
            }
        }
        return ben;
    }
    
    public static Builder lootingMultiplier(final RandomValueBounds csf) {
        return new Builder(csf);
    }
    
    public static class Builder extends LootItemConditionalFunction.Builder<Builder> {
        private final RandomValueBounds count;
        private int limit;
        
        public Builder(final RandomValueBounds csf) {
            this.limit = 0;
            this.count = csf;
        }
        
        @Override
        protected Builder getThis() {
            return this;
        }
        
        public Builder setLimit(final int integer) {
            this.limit = integer;
            return this;
        }
        
        @Override
        public LootItemFunction build() {
            return new LootingEnchantFunction(this.getConditions(), this.count, this.limit, null);
        }
    }
    
    public static class Serializer extends LootItemConditionalFunction.Serializer<LootingEnchantFunction> {
        protected Serializer() {
            super(new ResourceLocation("looting_enchant"), LootingEnchantFunction.class);
        }
        
        @Override
        public void serialize(final JsonObject jsonObject, final LootingEnchantFunction ctk, final JsonSerializationContext jsonSerializationContext) {
            super.serialize(jsonObject, ctk, jsonSerializationContext);
            jsonObject.add("count", jsonSerializationContext.serialize(ctk.value));
            if (ctk.hasLimit()) {
                jsonObject.add("limit", jsonSerializationContext.serialize(ctk.limit));
            }
        }
        
        @Override
        public LootingEnchantFunction deserialize(final JsonObject jsonObject, final JsonDeserializationContext jsonDeserializationContext, final LootItemCondition[] arr) {
            final int integer5 = GsonHelper.getAsInt(jsonObject, "limit", 0);
            return new LootingEnchantFunction(arr, GsonHelper.<RandomValueBounds>getAsObject(jsonObject, "count", jsonDeserializationContext, RandomValueBounds.class), integer5, null);
        }
    }
}
