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

import com.google.gson.JsonSyntaxException;
import net.minecraft.util.GsonHelper;
import com.google.common.collect.Lists;
import com.google.gson.JsonDeserializationContext;
import net.minecraft.resources.ResourceLocation;
import java.util.Iterator;
import com.google.gson.JsonElement;
import com.google.gson.JsonPrimitive;
import com.google.gson.JsonArray;
import com.google.gson.JsonSerializationContext;
import com.google.gson.JsonObject;
import com.google.common.collect.Sets;
import java.util.Set;
import org.apache.logging.log4j.LogManager;
import net.minecraft.world.item.EnchantedBookItem;
import net.minecraft.world.item.enchantment.EnchantmentInstance;
import net.minecraft.world.level.ItemLike;
import net.minecraft.util.Mth;
import java.util.Random;
import java.util.stream.Collector;
import java.util.stream.Collectors;
import net.minecraft.core.Registry;
import net.minecraft.world.item.Items;
import net.minecraft.world.level.storage.loot.LootContext;
import net.minecraft.world.item.ItemStack;
import com.google.common.collect.ImmutableList;
import java.util.Collection;
import net.minecraft.world.level.storage.loot.predicates.LootItemCondition;
import net.minecraft.world.item.enchantment.Enchantment;
import java.util.List;
import org.apache.logging.log4j.Logger;

public class EnchantRandomlyFunction extends LootItemConditionalFunction {
    private static final Logger LOGGER;
    private final List<Enchantment> enchantments;
    
    private EnchantRandomlyFunction(final LootItemCondition[] arr, final Collection<Enchantment> collection) {
        super(arr);
        this.enchantments = ImmutableList.copyOf(collection);
    }
    
    @Override
    public LootItemFunctionType getType() {
        return LootItemFunctions.ENCHANT_RANDOMLY;
    }
    
    public ItemStack run(final ItemStack bki, final LootContext dat) {
        final Random random5 = dat.getRandom();
        Enchantment bnw2;
        if (this.enchantments.isEmpty()) {
            final boolean boolean6 = bki.getItem() == Items.BOOK;
            final List<Enchantment> list7 = Registry.ENCHANTMENT.stream().filter(Enchantment::isDiscoverable).filter(bnw -> boolean6 || bnw.canEnchant(bki)).collect(Collectors.toList());
            if (list7.isEmpty()) {
                EnchantRandomlyFunction.LOGGER.warn("Couldn't find a compatible enchantment for {}", bki);
                return bki;
            }
            bnw2 = list7.get(random5.nextInt(list7.size()));
        }
        else {
            bnw2 = this.enchantments.get(random5.nextInt(this.enchantments.size()));
        }
        return enchantItem(bki, bnw2, random5);
    }
    
    private static ItemStack enchantItem(ItemStack bki, final Enchantment bnw, final Random random) {
        final int integer4 = Mth.nextInt(random, bnw.getMinLevel(), bnw.getMaxLevel());
        if (bki.getItem() == Items.BOOK) {
            bki = new ItemStack(Items.ENCHANTED_BOOK);
            EnchantedBookItem.addEnchantment(bki, new EnchantmentInstance(bnw, integer4));
        }
        else {
            bki.enchant(bnw, integer4);
        }
        return bki;
    }
    
    public static LootItemConditionalFunction.Builder<?> randomApplicableEnchantment() {
        return LootItemConditionalFunction.simpleBuilder(arr -> new EnchantRandomlyFunction(arr, ImmutableList.of()));
    }
    
    static {
        LOGGER = LogManager.getLogger();
    }
    
    public static class Builder extends LootItemConditionalFunction.Builder<Builder> {
        private final Set<Enchantment> enchantments;
        
        public Builder() {
            this.enchantments = Sets.newHashSet();
        }
        
        @Override
        protected Builder getThis() {
            return this;
        }
        
        public Builder withEnchantment(final Enchantment bnw) {
            this.enchantments.add(bnw);
            return this;
        }
        
        @Override
        public LootItemFunction build() {
            return new EnchantRandomlyFunction(this.getConditions(), this.enchantments, null);
        }
    }
    
    public static class Serializer extends LootItemConditionalFunction.Serializer<EnchantRandomlyFunction> {
        @Override
        public void serialize(final JsonObject jsonObject, final EnchantRandomlyFunction dca, final JsonSerializationContext jsonSerializationContext) {
            super.serialize(jsonObject, dca, jsonSerializationContext);
            if (!dca.enchantments.isEmpty()) {
                final JsonArray jsonArray5 = new JsonArray();
                for (final Enchantment bnw7 : dca.enchantments) {
                    final ResourceLocation uh8 = Registry.ENCHANTMENT.getKey(bnw7);
                    if (uh8 == null) {
                        throw new IllegalArgumentException("Don't know how to serialize enchantment " + bnw7);
                    }
                    jsonArray5.add(new JsonPrimitive(uh8.toString()));
                }
                jsonObject.add("enchantments", jsonArray5);
            }
        }
        
        @Override
        public EnchantRandomlyFunction deserialize(final JsonObject jsonObject, final JsonDeserializationContext jsonDeserializationContext, final LootItemCondition[] arr) {
            final List<Enchantment> list5 = Lists.newArrayList();
            if (jsonObject.has("enchantments")) {
                final JsonArray jsonArray6 = GsonHelper.getAsJsonArray(jsonObject, "enchantments");
                for (final JsonElement jsonElement8 : jsonArray6) {
                    final String string9 = GsonHelper.convertToString(jsonElement8, "enchantment");
                    final Object o;
                    final String str;
                    final Enchantment bnw10 = Registry.ENCHANTMENT.getOptional(new ResourceLocation(string9)).<Throwable>orElseThrow(() -> {
                        new JsonSyntaxException("Unknown enchantment '" + str + "'");
                        return o;
                    });
                    list5.add(bnw10);
                }
            }
            return new EnchantRandomlyFunction(arr, list5, null);
        }
    }
}
