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

import com.google.gson.JsonParseException;
import com.google.gson.JsonElement;
import net.minecraft.core.Registry;
import net.minecraft.util.GsonHelper;
import com.google.gson.JsonDeserializationContext;
import com.google.gson.JsonSerializationContext;
import com.google.gson.JsonObject;
import java.util.Random;
import com.google.common.collect.Maps;
import net.minecraft.world.item.enchantment.EnchantmentHelper;
import net.minecraft.world.level.storage.loot.LootContext;
import net.minecraft.world.item.ItemStack;
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.item.enchantment.Enchantment;
import net.minecraft.resources.ResourceLocation;
import java.util.Map;

public class ApplyBonusCount extends LootItemConditionalFunction {
    private static final Map<ResourceLocation, FormulaDeserializer> FORMULAS;
    private final Enchantment enchantment;
    private final Formula formula;
    
    private ApplyBonusCount(final LootItemCondition[] arr, final Enchantment bnw, final Formula b) {
        super(arr);
        this.enchantment = bnw;
        this.formula = b;
    }
    
    @Override
    public LootItemFunctionType getType() {
        return LootItemFunctions.APPLY_BONUS;
    }
    
    @Override
    public Set<LootContextParam<?>> getReferencedContextParams() {
        return ImmutableSet.of(LootContextParams.TOOL);
    }
    
    public ItemStack run(final ItemStack bki, final LootContext dat) {
        final ItemStack bki2 = dat.<ItemStack>getParamOrNull(LootContextParams.TOOL);
        if (bki2 != null) {
            final int integer5 = EnchantmentHelper.getItemEnchantmentLevel(this.enchantment, bki2);
            final int integer6 = this.formula.calculateNewCount(dat.getRandom(), bki.getCount(), integer5);
            bki.setCount(integer6);
        }
        return bki;
    }
    
    public static Builder<?> addBonusBinomialDistributionCount(final Enchantment bnw, final float float2, final int integer) {
        final ApplyBonusCount applyBonusCount;
        return LootItemConditionalFunction.simpleBuilder(arr -> {
            new ApplyBonusCount(arr, bnw, new BinomialWithBonusCount(integer, float2));
            return applyBonusCount;
        });
    }
    
    public static Builder<?> addOreBonusCount(final Enchantment bnw) {
        final ApplyBonusCount applyBonusCount;
        return LootItemConditionalFunction.simpleBuilder(arr -> {
            new ApplyBonusCount(arr, bnw, new OreDrops());
            return applyBonusCount;
        });
    }
    
    public static Builder<?> addUniformBonusCount(final Enchantment bnw) {
        final ApplyBonusCount applyBonusCount;
        return LootItemConditionalFunction.simpleBuilder(arr -> {
            new ApplyBonusCount(arr, bnw, new UniformBonusCount(1));
            return applyBonusCount;
        });
    }
    
    public static Builder<?> addUniformBonusCount(final Enchantment bnw, final int integer) {
        final ApplyBonusCount applyBonusCount;
        return LootItemConditionalFunction.simpleBuilder(arr -> {
            new ApplyBonusCount(arr, bnw, new UniformBonusCount(integer));
            return applyBonusCount;
        });
    }
    
    static {
        (FORMULAS = Maps.newHashMap()).put(BinomialWithBonusCount.TYPE, BinomialWithBonusCount::deserialize);
        ApplyBonusCount.FORMULAS.put(OreDrops.TYPE, OreDrops::deserialize);
        ApplyBonusCount.FORMULAS.put(UniformBonusCount.TYPE, UniformBonusCount::deserialize);
    }
    
    static final class BinomialWithBonusCount implements Formula {
        public static final ResourceLocation TYPE;
        private final int extraRounds;
        private final float probability;
        
        public BinomialWithBonusCount(final int integer, final float float2) {
            this.extraRounds = integer;
            this.probability = float2;
        }
        
        @Override
        public int calculateNewCount(final Random random, int integer2, final int integer3) {
            for (int integer4 = 0; integer4 < integer3 + this.extraRounds; ++integer4) {
                if (random.nextFloat() < this.probability) {
                    ++integer2;
                }
            }
            return integer2;
        }
        
        @Override
        public void serializeParams(final JsonObject jsonObject, final JsonSerializationContext jsonSerializationContext) {
            jsonObject.addProperty("extra", this.extraRounds);
            jsonObject.addProperty("probability", this.probability);
        }
        
        public static Formula deserialize(final JsonObject jsonObject, final JsonDeserializationContext jsonDeserializationContext) {
            final int integer3 = GsonHelper.getAsInt(jsonObject, "extra");
            final float float4 = GsonHelper.getAsFloat(jsonObject, "probability");
            return new BinomialWithBonusCount(integer3, float4);
        }
        
        @Override
        public ResourceLocation getType() {
            return BinomialWithBonusCount.TYPE;
        }
        
        static {
            TYPE = new ResourceLocation("binomial_with_bonus_count");
        }
    }
    
    static final class UniformBonusCount implements Formula {
        public static final ResourceLocation TYPE;
        private final int bonusMultiplier;
        
        public UniformBonusCount(final int integer) {
            this.bonusMultiplier = integer;
        }
        
        @Override
        public int calculateNewCount(final Random random, final int integer2, final int integer3) {
            return integer2 + random.nextInt(this.bonusMultiplier * integer3 + 1);
        }
        
        @Override
        public void serializeParams(final JsonObject jsonObject, final JsonSerializationContext jsonSerializationContext) {
            jsonObject.addProperty("bonusMultiplier", this.bonusMultiplier);
        }
        
        public static Formula deserialize(final JsonObject jsonObject, final JsonDeserializationContext jsonDeserializationContext) {
            final int integer3 = GsonHelper.getAsInt(jsonObject, "bonusMultiplier");
            return new UniformBonusCount(integer3);
        }
        
        @Override
        public ResourceLocation getType() {
            return UniformBonusCount.TYPE;
        }
        
        static {
            TYPE = new ResourceLocation("uniform_bonus_count");
        }
    }
    
    static final class OreDrops implements Formula {
        public static final ResourceLocation TYPE;
        
        private OreDrops() {
        }
        
        @Override
        public int calculateNewCount(final Random random, final int integer2, final int integer3) {
            if (integer3 > 0) {
                int integer4 = random.nextInt(integer3 + 2) - 1;
                if (integer4 < 0) {
                    integer4 = 0;
                }
                return integer2 * (integer4 + 1);
            }
            return integer2;
        }
        
        @Override
        public void serializeParams(final JsonObject jsonObject, final JsonSerializationContext jsonSerializationContext) {
        }
        
        public static Formula deserialize(final JsonObject jsonObject, final JsonDeserializationContext jsonDeserializationContext) {
            return new OreDrops();
        }
        
        @Override
        public ResourceLocation getType() {
            return OreDrops.TYPE;
        }
        
        static {
            TYPE = new ResourceLocation("ore_drops");
        }
    }
    
    public static class Serializer extends LootItemConditionalFunction.Serializer<ApplyBonusCount> {
        @Override
        public void serialize(final JsonObject jsonObject, final ApplyBonusCount dbv, final JsonSerializationContext jsonSerializationContext) {
            super.serialize(jsonObject, dbv, jsonSerializationContext);
            jsonObject.addProperty("enchantment", Registry.ENCHANTMENT.getKey(dbv.enchantment).toString());
            jsonObject.addProperty("formula", dbv.formula.getType().toString());
            final JsonObject jsonObject2 = new JsonObject();
            dbv.formula.serializeParams(jsonObject2, jsonSerializationContext);
            if (jsonObject2.size() > 0) {
                jsonObject.add("parameters", jsonObject2);
            }
        }
        
        @Override
        public ApplyBonusCount deserialize(final JsonObject jsonObject, final JsonDeserializationContext jsonDeserializationContext, final LootItemCondition[] arr) {
            final ResourceLocation uh5 = new ResourceLocation(GsonHelper.getAsString(jsonObject, "enchantment"));
            final Object o;
            final Object obj;
            final Enchantment bnw6 = Registry.ENCHANTMENT.getOptional(uh5).<Throwable>orElseThrow(() -> {
                new JsonParseException("Invalid enchantment id: " + obj);
                return o;
            });
            final ResourceLocation uh6 = new ResourceLocation(GsonHelper.getAsString(jsonObject, "formula"));
            final FormulaDeserializer c8 = ApplyBonusCount.FORMULAS.get(uh6);
            if (c8 == null) {
                throw new JsonParseException("Invalid formula id: " + uh6);
            }
            Formula b9;
            if (jsonObject.has("parameters")) {
                b9 = c8.deserialize(GsonHelper.getAsJsonObject(jsonObject, "parameters"), jsonDeserializationContext);
            }
            else {
                b9 = c8.deserialize(new JsonObject(), jsonDeserializationContext);
            }
            return new ApplyBonusCount(arr, bnw6, b9, null);
        }
    }
    
    interface Formula {
        int calculateNewCount(final Random random, final int integer2, final int integer3);
        
        void serializeParams(final JsonObject jsonObject, final JsonSerializationContext jsonSerializationContext);
        
        ResourceLocation getType();
    }
    
    interface FormulaDeserializer {
        Formula deserialize(final JsonObject jsonObject, final JsonDeserializationContext jsonDeserializationContext);
    }
}
