package net.minestom.vanilla.loot;

import net.kyori.adventure.key.Key;
import net.kyori.adventure.nbt.*;
import net.kyori.adventure.text.Component;
import net.kyori.adventure.util.RGBLike;
import net.minestom.server.MinecraftServer;
import net.minestom.server.ServerFlag;
import net.minestom.server.codec.Codec;
import net.minestom.server.codec.Result;
import net.minestom.server.codec.StructCodec;
import net.minestom.server.codec.Transcoder;
import net.minestom.server.color.Color;
import net.minestom.server.component.DataComponent;
import net.minestom.server.component.DataComponentMap;
import net.minestom.server.component.DataComponents;
import net.minestom.server.entity.*;
import net.minestom.server.entity.attribute.Attribute;
import net.minestom.server.entity.attribute.AttributeModifier;
import net.minestom.server.entity.attribute.AttributeOperation;
import net.minestom.server.instance.block.Block;
import net.minestom.server.item.ItemStack;
import net.minestom.server.item.Material;
import net.minestom.server.item.book.FilteredText;
import net.minestom.server.item.component.*;
import net.minestom.server.item.enchant.Enchantment;
import net.minestom.server.item.instrument.Instrument;
import net.minestom.server.potion.PotionEffect;
import net.minestom.server.potion.PotionType;
import net.minestom.server.registry.DynamicRegistry;
import net.minestom.server.registry.Registries;
import net.minestom.server.registry.RegistryKey;
import net.minestom.server.registry.RegistryTag;
import net.minestom.server.tag.Tag;
import net.minestom.server.utils.Either;
import net.minestom.vanilla.loot.util.*;
import net.minestom.vanilla.loot.util.nbt.NBTPath;
import net.minestom.vanilla.loot.util.nbt.NBTReference;
import net.minestom.vanilla.loot.util.nbt.NBTUtils;
import net.minestom.vanilla.loot.util.predicate.ItemPredicate;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;

import java.util.*;
import java.util.function.Function;
import java.util.stream.Collectors;
import java.util.stream.Stream;

/**
 * A function that allows loot to pass through it, potentially making modifications.
 */
@SuppressWarnings("UnstableApiUsage")
public interface LootFunction {

    @NotNull Codec<LootFunction> CODEC = makeCodec().orElse(makeCodec().list().transform(Sequence::new, seq -> ((Sequence) seq).functions()));

    private static StructCodec<LootFunction> makeCodec() {
        return Codec.RegistryTaggedUnion(registries -> {
            class Holder {
                static final @NotNull DynamicRegistry<StructCodec<? extends LootFunction>> CODEC = createDefaultRegistry();
            }
            return Holder.CODEC;
        }, LootFunction::codec, "function");
    }

    static @NotNull DynamicRegistry<StructCodec<? extends LootFunction>> createDefaultRegistry() {
        final DynamicRegistry<StructCodec<? extends LootFunction>> registry = DynamicRegistry.create(Key.key("loot_functions"));
        registry.register("apply_bonus", ApplyBonus.CODEC);
        registry.register("copy_components", CopyComponents.CODEC);
        registry.register("copy_custom_data", CopyCustomData.CODEC);
        registry.register("copy_name", CopyName.CODEC);
        registry.register("copy_state", CopyState.CODEC);
        registry.register("enchanted_count_increase", EnchantedCountIncrease.CODEC);
        registry.register("enchant_randomly", EnchantRandomly.CODEC);
        registry.register("enchant_with_levels", EnchantWithLevels.CODEC);
        registry.register("exploration_map", ExplorationMap.CODEC);
        registry.register("explosion_decay", ExplosionDecay.CODEC);
        registry.register("fill_player_head", FillPlayerHead.CODEC);
        registry.register("filtered", Filtered.CODEC);
        registry.register("furnace_smelt", FurnaceSmelt.CODEC);
        registry.register("limit_count", LimitCount.CODEC);
        registry.register("modify_contents", ModifyContents.CODEC);
        registry.register("reference", Reference.CODEC);
        registry.register("set_attributes", SetAttributes.CODEC);
        registry.register("set_banner_pattern", SetBannerPattern.CODEC);
        registry.register("set_book_cover", SetBookCover.CODEC);
        registry.register("set_components", SetComponents.CODEC);
        registry.register("set_contents", SetContents.CODEC);
        registry.register("set_count", SetCount.CODEC);
        registry.register("set_custom_data", SetCustomData.CODEC);
        registry.register("set_custom_model_data", SetCustomModelData.CODEC);
        registry.register("set_damage", SetDamage.CODEC);
        registry.register("set_enchantments", SetEnchantments.CODEC);
        registry.register("set_firework_explosion", SetFireworkExplosion.CODEC);
        registry.register("set_fireworks", SetFireworks.CODEC);
        registry.register("set_instrument", SetInstrument.CODEC);
        registry.register("set_item", SetItem.CODEC);
        registry.register("set_loot_table", SetLootTable.CODEC);
        registry.register("set_lore", SetLore.CODEC);
        registry.register("set_name", SetName.CODEC);
        registry.register("set_ominous_bottle_amplifier", SetOminousBottleAmplifier.CODEC);
        registry.register("set_potion", SetPotion.CODEC);
        registry.register("set_stew_effect", SetStewEffect.CODEC);
        registry.register("set_writable_book_pages", SetWritableBookPages.CODEC);
        registry.register("set_written_book_pages", SetWrittenBookPages.CODEC);
        registry.register("sequence", Sequence.CODEC);
        registry.register("toggle_tooltips", ToggleTooltips.CODEC);
        return registry;
    }

    /**
     * Performs any mutations on the provided object and returns the result.
     * @param input the input item to this function
     * @param context the context object, to use if required
     * @return the modified form of the input
     */
    @NotNull ItemStack apply(@NotNull ItemStack input, @NotNull LootContext context);

    /**
     * @return the codec that can encode this function
     */
    @NotNull StructCodec<? extends LootFunction> codec();

    /**
     * Applies each function to the given item consecutively.
     * @param functions the functions to apply
     * @param item the item to modify
     * @param context the context to use
     * @return the modified item
     */
    static @NotNull ItemStack apply(@NotNull Collection<LootFunction> functions, @NotNull ItemStack item, @NotNull LootContext context) {
        for (LootFunction function : functions) {
            item = function.apply(item, context);
        }
        return item;
    }

    /**
     * Applies each function to each of the given items consecutively.
     * @param functions the functions to apply
     * @param items the items to modify
     * @param context the context to use
     * @return the modified items
     */
    static @NotNull List<ItemStack> apply(@NotNull Collection<LootFunction> functions, @NotNull List<ItemStack> items, @NotNull LootContext context) {
        List<ItemStack> newItems = new ArrayList<>(items.size());
        for (ItemStack item : items) {
            newItems.add(LootFunction.apply(functions, item, context));
        }
        return newItems;
    }

    record ApplyBonus(@NotNull List<LootPredicate> predicates, @NotNull RegistryKey<Enchantment> enchantment, @NotNull Formula.Wrapper formula) implements LootFunction {
        public static final @NotNull StructCodec<ApplyBonus> CODEC = StructCodec.struct(
                "conditions", LootPredicate.CODEC.list().optional(List.of()), ApplyBonus::predicates,
                "enchantment", RegistryKey.codec(Registries::enchantment), ApplyBonus::enchantment,
                StructCodec.INLINE, Formula.CODEC, ApplyBonus::formula,
                ApplyBonus::new
        );

        public sealed interface Formula {

            enum FormulaType {
                BINOMIAL_WITH_BONUS_COUNT,
                ORE_DROPS,
                UNIFORM_BONUS_COUNT,
            }

            record Wrapper(Formula parameters) {
                @SuppressWarnings("unchecked")
                private static <T extends Formula> StructCodec<Wrapper> wrap(@NotNull StructCodec<T> codec) {
                    return StructCodec.struct(
                            "parameters", codec.transform(a->a,a->(T)a), Wrapper::parameters,
                            Wrapper::new
                    );
                }

                private static final StructCodec<Wrapper> BINOMIAL_WITH_BONUS_COUNT = wrap(BinomialWithBonusCount.CODEC);
                private static final StructCodec<Wrapper> ORE_DROPS = new StructCodec<>() {
                    @Override
                    public @NotNull <D> Result<Wrapper> decodeFromMap(@NotNull Transcoder<D> coder, Transcoder.@NotNull MapLike<D> map) {
                        return new Result.Ok<>(new Wrapper(new OreDrops()));
                    }

                    @Override
                    public @NotNull <D> Result<D> encodeToMap(@NotNull Transcoder<D> coder, @NotNull Wrapper value, Transcoder.@NotNull MapBuilder<D> map) {
                        return new Result.Ok<>(map.build());
                    }
                };
                private static final StructCodec<Wrapper> UNIFORM_BONUS_COUNT = wrap(UniformBonusCount.CODEC);

                public static @NotNull StructCodec<Wrapper> codec(@NotNull FormulaType type) {
                    return switch (type) {
                        case BINOMIAL_WITH_BONUS_COUNT -> BINOMIAL_WITH_BONUS_COUNT;
                        case ORE_DROPS -> ORE_DROPS;
                        case UNIFORM_BONUS_COUNT -> UNIFORM_BONUS_COUNT;
                    };
                }
            }

            @NotNull Codec<Wrapper> CODEC = Codec.KEY.transform(key -> switch (key.asString()) {
                case "minecraft:binomial_with_bonus_count" -> FormulaType.BINOMIAL_WITH_BONUS_COUNT;
                case "minecraft:ore_drops" -> FormulaType.ORE_DROPS;
                case "minecraft:uniform_bonus_count" -> FormulaType.UNIFORM_BONUS_COUNT;
                default -> throw new IllegalArgumentException();
            }, type -> Key.key(type.toString().toLowerCase())).unionType("formula", Wrapper::codec, (Wrapper formula) -> switch (formula.parameters()) {
                case BinomialWithBonusCount ignored -> FormulaType.BINOMIAL_WITH_BONUS_COUNT;
                case OreDrops ignored -> FormulaType.ORE_DROPS;
                case UniformBonusCount ignored -> FormulaType.UNIFORM_BONUS_COUNT;
            });

            int calculate(@NotNull Random random, int count, int level);

            record UniformBonusCount(int bonusMultiplier) implements Formula {
                public static final @NotNull StructCodec<UniformBonusCount> CODEC = StructCodec.struct(
                        "bonusMultiplier", StructCodec.INT, UniformBonusCount::bonusMultiplier,
                        UniformBonusCount::new
                );

                @Override
                public int calculate(@NotNull Random random, int count, int level) {
                    return count + random.nextInt(bonusMultiplier * level + 1);
                }
            }

            record OreDrops() implements Formula {
                public static final @NotNull StructCodec<OreDrops> CODEC = StructCodec.struct(OreDrops::new);

                @Override
                public int calculate(@NotNull Random random, int count, int level) {
                    if (level <= 0) return count;

                    return count * Math.max(1, random.nextInt(level + 2));
                }
            }

            record BinomialWithBonusCount(float probability, int extra) implements Formula {
                public static final @NotNull StructCodec<BinomialWithBonusCount> CODEC = StructCodec.struct(
                        "probability", Codec.FLOAT, BinomialWithBonusCount::probability,
                        "extra", Codec.INT, BinomialWithBonusCount::extra,
                        BinomialWithBonusCount::new
                );

                @Override
                public int calculate(@NotNull Random random, int count, int level) {
                    for (int i = 0; i < extra + level; i++) {
                        if (random.nextFloat() < probability) {
                            count++;
                        }
                    }

                    return count;
                }
            }

        }

        @Override
        public @NotNull ItemStack apply(@NotNull ItemStack input, @NotNull LootContext context) {
            ItemStack tool = context.get(LootContext.TOOL);
            if (tool == null) return input;

            int level = EnchantmentUtils.level(tool, enchantment);
            int newCount = formula.parameters().calculate(context.require(LootContext.RANDOM), input.amount(), level);

            return input.withAmount(newCount);
        }

        @Override
        public @NotNull StructCodec<? extends LootFunction> codec() {
            return CODEC;
        }
    }

    record CopyComponents(@NotNull List<LootPredicate> predicates, @NotNull RelevantTarget source,
                          @Nullable List<DataComponent<?>> include, @Nullable List<DataComponent<?>> exclude) implements LootFunction {
        public static final @NotNull StructCodec<CopyComponents> CODEC = StructCodec.struct(
                "conditions", LootPredicate.CODEC.list().optional(List.of()), CopyComponents::predicates,
                "source", RelevantTarget.CODEC, CopyComponents::source,
                "include", Codec.KEY.<DataComponent<?>>transform(DataComponent::fromKey, DataComponent::key).list().optional(), CopyComponents::include,
                "exclude", Codec.KEY.<DataComponent<?>>transform(DataComponent::fromKey, DataComponent::key).list().optional(), CopyComponents::exclude,
                CopyComponents::new
        );

        @Override
        public @NotNull ItemStack apply(@NotNull ItemStack input, @NotNull LootContext context) {
            if (!LootPredicate.all(predicates, context)) return input;

            // TODO: Incomplete
            throw new UnsupportedOperationException("TODO: Implement Tag<DataComponentMap> for blocks.");
        }

        @Override
        public @NotNull StructCodec<? extends LootFunction> codec() {
            return CODEC;
        }
    }
    
    record CopyCustomData(@NotNull List<LootPredicate> predicates, @NotNull LootNBT source, @NotNull List<Operation> ops) implements LootFunction {
        public static final @NotNull StructCodec<CopyCustomData> CODEC = StructCodec.struct(
                "conditions", LootPredicate.CODEC.list().optional(List.of()), CopyCustomData::predicates,
                "source", LootNBT.CODEC, CopyCustomData::source,
                "ops", Operation.CODEC.list(), CopyCustomData::ops,
                CopyCustomData::new
        );

        public record Operation(@NotNull NBTPath source, @NotNull NBTPath target, @NotNull Operator op) {
            public static final @NotNull StructCodec<Operation> CODEC = StructCodec.struct(
                    "source", NBTPath.CODEC, Operation::source,
                    "target", NBTPath.CODEC, Operation::target,
                    "op", Operator.SERIALIZER, Operation::op,
                    Operation::new
            );

            public void execute(@NotNull NBTReference nbt, @NotNull BinaryTag sourceTag) {
                List<BinaryTag> nbts = new ArrayList<>();
                source.get(sourceTag).forEach(ref -> nbts.add(ref.get()));

                if (nbts.isEmpty()) return;
                op.merge(nbt, target, nbts);
            }
        }

        public enum Operator {
            REPLACE() {
                @Override
                public void merge(@NotNull NBTReference nbt, @NotNull NBTPath target, @NotNull List<BinaryTag> source) {
                    target.set(nbt, source.getLast());
                }
            },
            APPEND() {
                @Override
                public void merge(@NotNull NBTReference nbt, @NotNull NBTPath target, @NotNull List<BinaryTag> source) {
                    List<NBTReference> nbts = target.getWithDefaults(nbt, ListBinaryTag::empty);

                    for (var ref : nbts) {
                        source.forEach(ref::listAdd);
                    }
                }
            },
            MERGE() {
                @Override
                public void merge(@NotNull NBTReference nbt, @NotNull NBTPath target, @NotNull List<BinaryTag> source) {
                    List<NBTReference> nbts = target.getWithDefaults(nbt, CompoundBinaryTag::empty);

                    for (var ref : nbts) {
                        if (ref.get() instanceof CompoundBinaryTag compound) {
                            for (var nbt2 : source) {
                                if (nbt2 instanceof CompoundBinaryTag compound2) {
                                    ref.set(NBTUtils.merge(compound, compound2));
                                }
                            }
                        }
                    }
                }
            };

            public static final @NotNull Codec<Operator> SERIALIZER = Codec.Enum(Operator.class);

            public abstract void merge(@NotNull NBTReference nbt, @NotNull NBTPath target, @NotNull List<BinaryTag> source);
        }

        @Override
        public @NotNull ItemStack apply(@NotNull ItemStack input, @NotNull LootContext context) {
            if (!LootPredicate.all(predicates, context)) return input;

            BinaryTag sourceNBT = source.getNBT(context);
            if (sourceNBT == null) return input;

            NBTReference targetNBT = NBTReference.of(input.get(DataComponents.CUSTOM_DATA, CustomData.EMPTY).nbt());

            for (Operation operation : ops) {
                operation.execute(targetNBT, sourceNBT);
            }

            if (targetNBT.get() instanceof CompoundBinaryTag compound) {
                return input.with(DataComponents.CUSTOM_DATA, new CustomData(compound));
            } else {
                return input;
            }
        }

        @Override
        public @NotNull StructCodec<? extends LootFunction> codec() {
            return CODEC;
        }
    }
    
    record CopyName(@NotNull List<LootPredicate> predicates, @NotNull RelevantTarget source) implements LootFunction {
        public static final @NotNull StructCodec<CopyName> CODEC = StructCodec.struct(
                "conditions", LootPredicate.CODEC.list().optional(List.of()), CopyName::predicates,
                "source", RelevantTarget.CODEC, CopyName::source,
                CopyName::new
        );

        private static final Tag<Component> CUSTOM_NAME = Tag.Component("CustomName");

        @Override
        public @NotNull ItemStack apply(@NotNull ItemStack input, @NotNull LootContext context) {
            if (!LootPredicate.all(predicates, context)) return input;

            Object key = context.get(source.key());

            Component customName;
            if (key instanceof Entity entity && entity.getCustomName() != null) {
                customName = entity.getCustomName();
            } else if (key instanceof Block block && block.hasTag(CUSTOM_NAME)) {
                customName = block.getTag(CUSTOM_NAME);
            } else {
                return input;
            }

            return input.with(DataComponents.CUSTOM_NAME, customName);
        }

        @Override
        public @NotNull StructCodec<? extends LootFunction> codec() {
            return CODEC;
        }
    }

    record CopyState(@NotNull List<LootPredicate> predicates, @NotNull Block block, @NotNull List<String> properties) implements LootFunction {
        public static final @NotNull StructCodec<CopyState> CODEC = StructCodec.struct(
                "conditions", LootPredicate.CODEC.list().optional(List.of()), CopyState::predicates,
                "block", Codec.KEY.transform(Block::fromKey, Block::key), CopyState::block,
                "properties", Codec.STRING.list(), CopyState::properties,
                CopyState::new
        );

        public CopyState {
            List<String> props = new ArrayList<>(properties);
            props.removeIf(name -> !block.properties().containsKey(name));
            properties = List.copyOf(props);
        }

        @Override
        public @NotNull ItemStack apply(@NotNull ItemStack input, @NotNull LootContext context) {
            if (!LootPredicate.all(predicates, context)) return input;

            Block block = context.get(LootContext.BLOCK_STATE);
            if (block == null) return input;

            ItemBlockState irritableBowelSyndrome = input.get(DataComponents.BLOCK_STATE, ItemBlockState.EMPTY);

            if (!block.key().equals(this.block.key())) return input;

            for (var prop : properties) {
                @Nullable String value = block.getProperty(prop);
                if (value == null) continue;

                irritableBowelSyndrome = irritableBowelSyndrome.with(prop, value);
            }

            return input.with(DataComponents.BLOCK_STATE, irritableBowelSyndrome);
        }

        @Override
        public @NotNull StructCodec<? extends LootFunction> codec() {
            return CODEC;
        }
    }

    record EnchantedCountIncrease(@NotNull List<LootPredicate> predicates, @NotNull RegistryKey<Enchantment> enchantment,
                                  @NotNull LootNumber count, @Nullable Integer limit) implements LootFunction {
        public static final @NotNull StructCodec<EnchantedCountIncrease> CODEC = StructCodec.struct(
                "conditions", LootPredicate.CODEC.list().optional(List.of()), EnchantedCountIncrease::predicates,
                "enchantment", RegistryKey.codec(Registries::enchantment), EnchantedCountIncrease::enchantment,
                "count", LootNumber.CODEC, EnchantedCountIncrease::count,
                "limit", Codec.INT.optional(), EnchantedCountIncrease::limit,
                EnchantedCountIncrease::new
        );

        @Override
        public @NotNull ItemStack apply(@NotNull ItemStack input, @NotNull LootContext context) {
            if (!LootPredicate.all(predicates, context)) return input;

            Entity attacker = context.get(LootContext.ATTACKING_ENTITY);
            int level = EnchantmentUtils.level(attacker, enchantment);

            if (level == 0) return input;

            int newAmount = input.amount() + level * count.getInt(context);

            return input.withAmount(limit != null ? Math.min(limit, newAmount) : newAmount);
        }

        @Override
        public @NotNull StructCodec<? extends LootFunction> codec() {
            return CODEC;
        }
    }

    record EnchantRandomly(@NotNull List<LootPredicate> predicates, @Nullable RegistryTag<Enchantment> options, boolean onlyCompatible) implements LootFunction {
        public static final @NotNull StructCodec<EnchantRandomly> CODEC = StructCodec.struct(
                "conditions", LootPredicate.CODEC.list().optional(List.of()), EnchantRandomly::predicates,
                "options", RegistryTag.codec(Registries::enchantment).optional(), EnchantRandomly::options,
                "only_compatible", Codec.BOOLEAN.optional(true), EnchantRandomly::onlyCompatible,
                EnchantRandomly::new
        );

        @Override
        public @NotNull ItemStack apply(@NotNull ItemStack input, @NotNull LootContext context) {
            var reg = MinecraftServer.getEnchantmentRegistry();

            List<RegistryKey<Enchantment>> values = new ArrayList<>();
            (options == null ? reg.keys() : options).forEach(values::add);

            if (onlyCompatible && !input.material().equals(Material.BOOK)) {
                values.removeIf(ench -> !reg.get(ench).supportedItems().contains(input.material()));
            }

            if (values.isEmpty()) return input;

            Random rng = context.require(LootContext.RANDOM);

            RegistryKey<Enchantment> chosen = values.get(rng.nextInt(values.size()));

            int level = rng.nextInt(reg.get(chosen).maxLevel() + 1);

            return EnchantmentUtils.modifyItem(input, map -> map.put(chosen, level));
        }

        @Override
        public @NotNull StructCodec<? extends LootFunction> codec() {
            return CODEC;
        }
    }

    record EnchantWithLevels(@NotNull List<LootPredicate> predicates, @NotNull LootNumber levels, @Nullable RegistryTag<Enchantment> options) implements LootFunction {
        public static final @NotNull StructCodec<EnchantWithLevels> CODEC = StructCodec.struct(
                "conditions", LootPredicate.CODEC.list().optional(List.of()), EnchantWithLevels::predicates,
                "levels", LootNumber.CODEC, EnchantWithLevels::levels,
                "options", RegistryTag.codec(Registries::enchantment).optional(), EnchantWithLevels::options,
                EnchantWithLevels::new
        );

        @Override
        public @NotNull ItemStack apply(@NotNull ItemStack input, @NotNull LootContext context) {
            if (!LootPredicate.all(predicates, context)) return input;

            if (true) throw new UnsupportedOperationException("TODO: Implement enchanting (Random, ItemStack, int levels, @Nullable RegistryTag<Enchantment> options -> ItemStack)");
            return null;
        }

        @Override
        public @NotNull StructCodec<? extends LootFunction> codec() {
            return CODEC;
        }
    }

    record ExplorationMap(@NotNull List<LootPredicate> predicates) implements LootFunction {
        public static final @NotNull StructCodec<ExplorationMap> CODEC = StructCodec.struct(
                "conditions", LootPredicate.CODEC.list().optional(List.of()), ExplorationMap::predicates,
                ExplorationMap::new
        );

        @Override
        public @NotNull ItemStack apply(@NotNull ItemStack input, @NotNull LootContext context) {
            if (!LootPredicate.all(predicates, context)) return input;

            // TODO: Incomplete
            throw new UnsupportedOperationException("TODO: Implement ExplorationMap functionality and serialization");
        }

        @Override
        public @NotNull StructCodec<? extends LootFunction> codec() {
            return CODEC;
        }
    }

    record ExplosionDecay(@NotNull List<LootPredicate> predicates) implements LootFunction {
        public static final @NotNull StructCodec<ExplosionDecay> CODEC = StructCodec.struct(
                "conditions", LootPredicate.CODEC.list().optional(List.of()), ExplosionDecay::predicates,
                ExplosionDecay::new
        );

        @Override
        public @NotNull ItemStack apply(@NotNull ItemStack input, @NotNull LootContext context) {
            if (!LootPredicate.all(predicates, context)) return input;

            Float radius = context.get(LootContext.EXPLOSION_RADIUS);
            if (radius == null) return input;

            Random random = context.require(LootContext.RANDOM);

            float chance = 1 / radius;
            int trials = input.amount();

            int newAmount = 0;

            for (int i = 0; i < trials; i++) {
                if (random.nextFloat() <= chance) {
                    newAmount++;
                }
            }

            return input.withAmount(newAmount);
        }

        @Override
        public @NotNull StructCodec<? extends LootFunction> codec() {
            return CODEC;
        }
    }

    record FillPlayerHead(@NotNull List<LootPredicate> predicates, @NotNull RelevantEntity entity) implements LootFunction {
        public static final @NotNull StructCodec<FillPlayerHead> CODEC = StructCodec.struct(
                "conditions", LootPredicate.CODEC.list().optional(List.of()), FillPlayerHead::predicates,
                "entity", RelevantEntity.CODEC, FillPlayerHead::entity,
                FillPlayerHead::new
        );

        @Override
        public @NotNull ItemStack apply(@NotNull ItemStack input, @NotNull LootContext context) {
            if (!LootPredicate.all(predicates, context)) return input;

            if (!input.material().equals(Material.PLAYER_HEAD)) return input;

            if (!(context.get(entity.key()) instanceof Player player)) return input;

            PlayerSkin skin = player.getSkin();
            if (skin == null) return input;

            return input.with(DataComponents.PROFILE, new HeadProfile(skin));
        }

        @Override
        public @NotNull StructCodec<? extends LootFunction> codec() {
            return CODEC;
        }
    }
    
    record Filtered(@NotNull List<LootPredicate> predicates, @NotNull ItemPredicate predicate, @NotNull LootFunction modifier) implements LootFunction {
        public static final @NotNull StructCodec<Filtered> CODEC = StructCodec.struct(
                "conditions", LootPredicate.CODEC.list().optional(List.of()), Filtered::predicates,
                "item_filter", ItemPredicate.CODEC, Filtered::predicate,
                "modifier", LootFunction.CODEC, Filtered::modifier,
                Filtered::new
        );

        @Override
        public @NotNull ItemStack apply(@NotNull ItemStack input, @NotNull LootContext context) {
            return LootPredicate.all(predicates, context) && predicate.test(input, context) ?
                    modifier.apply(input, context) : input;
        }

        @Override
        public @NotNull StructCodec<? extends LootFunction> codec() {
            return CODEC;
        }
    }
    
    record FurnaceSmelt(@NotNull List<LootPredicate> predicates) implements LootFunction {
        public static final @NotNull StructCodec<FurnaceSmelt> CODEC = StructCodec.struct(
                "conditions", LootPredicate.CODEC.list().optional(List.of()), FurnaceSmelt::predicates,
                FurnaceSmelt::new
        );

        @Override
        public @NotNull ItemStack apply(@NotNull ItemStack input, @NotNull LootContext context) {
            if (!LootPredicate.all(predicates, context)) return input;

            if (true) throw new UnsupportedOperationException("TODO: Implement smelting (ItemStack -> ItemStack)");
            ItemStack smelted = null;

            return smelted != null ? smelted.withAmount(input.amount()) : input;
        }

        @Override
        public @NotNull StructCodec<? extends LootFunction> codec() {
            return CODEC;
        }
    }

    record LimitCount(@NotNull List<LootPredicate> predicates, @NotNull LootNumberRange limit) implements LootFunction {
        public static final @NotNull StructCodec<LimitCount> CODEC = StructCodec.struct(
                "conditions", LootPredicate.CODEC.list().optional(List.of()), LimitCount::predicates,
                "limit", LootNumberRange.CODEC, LimitCount::limit,
                LimitCount::new
        );

        @Override
        public @NotNull ItemStack apply(@NotNull ItemStack input, @NotNull LootContext context) {
            if (!LootPredicate.all(predicates, context)) return input;
            return input.withAmount(i -> (int) limit.limit(context, i));
        }

        @Override
        public @NotNull StructCodec<? extends LootFunction> codec() {
            return CODEC;
        }
    }

    record ModifyContents(@NotNull List<LootPredicate> predicates, @NotNull List<LootFunction> modifier,
                          @NotNull DataComponent<List<ItemStack>> component) implements LootFunction {
        private static final @NotNull Map<Key, DataComponent<List<ItemStack>>> NAMED_CONTAINERS =
                Stream.of(DataComponents.CONTAINER, DataComponents.BUNDLE_CONTENTS, DataComponents.CHARGED_PROJECTILES)
                        .collect(Collectors.toMap(DataComponent::key, Function.identity()));

        private static final @NotNull Codec<DataComponent<List<ItemStack>>> CONTAINER = Codec.KEY.transform(NAMED_CONTAINERS::get, DataComponent::key);

        public static final @NotNull StructCodec<ModifyContents> CODEC = StructCodec.struct(
                "conditions", LootPredicate.CODEC.list().optional(List.of()), ModifyContents::predicates,
                "modifier", LootFunction.CODEC.list(), ModifyContents::modifier,
                "component", CONTAINER, ModifyContents::component,
                ModifyContents::new
        );

        @Override
        public @NotNull ItemStack apply(@NotNull ItemStack input, @NotNull LootContext context) {
            if (!LootPredicate.all(predicates, context)) return input;

            List<ItemStack> items = input.get(component);
            if (items == null) return input;

