// SPDX-FileCopyrightText: 2020 The CC: Tweaked Developers
//
// SPDX-License-Identifier: MPL-2.0

package dan200.computercraft.data;

import com.mojang.serialization.Codec;
import net.fabricmc.fabric.api.client.datagen.v1.provider.FabricModelProvider;
import net.fabricmc.fabric.api.datagen.v1.DataGeneratorEntrypoint;
import net.fabricmc.fabric.api.datagen.v1.FabricDataGenerator;
import net.fabricmc.fabric.api.datagen.v1.FabricPackOutput;
import net.fabricmc.fabric.api.datagen.v1.provider.FabricCodecDataProvider;
import net.fabricmc.fabric.api.datagen.v1.provider.FabricDynamicRegistryProvider;
import net.fabricmc.fabric.api.datagen.v1.provider.FabricTagsProvider;
import net.fabricmc.fabric.api.event.registry.DynamicRegistries;
import net.minecraft.client.data.models.BlockModelGenerators;
import net.minecraft.client.data.models.ItemModelGenerators;
import net.minecraft.core.Holder;
import net.minecraft.core.HolderLookup;
import net.minecraft.core.RegistrySetBuilder;
import net.minecraft.data.DataProvider;
import net.minecraft.data.PackOutput;
import net.minecraft.data.tags.TagsProvider;
import net.minecraft.resources.Identifier;
import net.minecraft.world.item.Item;
import net.minecraft.world.level.block.Block;

import java.util.concurrent.CompletableFuture;
import java.util.function.BiConsumer;
import java.util.function.Consumer;

public class FabricDataProviders implements DataGeneratorEntrypoint {
    @Override
    public void onInitializeDataGenerator(FabricDataGenerator generator) {
        var pack = new GeneratorSinkImpl(generator.createPack(), generator.getRegistries());
        DataProviders.add(pack);
    }

    private record GeneratorSinkImpl(
        FabricDataGenerator.Pack generator, CompletableFuture<HolderLookup.Provider> registries
    ) implements DataProviders.GeneratorSink {
        private <T extends DataProvider> T addWithFabricOutput(FabricDataGenerator.Pack.Factory<T> factory) {
            return generator.addProvider((FabricPackOutput p) -> new PrettyDataProvider<>(factory.create(p))).provider();
        }

        private <T extends DataProvider> T addWithRegistries(FabricDataGenerator.Pack.RegistryDependentFactory<T> factory) {
            return generator.addProvider((r, p) -> new PrettyDataProvider<>(factory.create(r, p))).provider();
        }

        @Override
        public <T extends DataProvider> T add(DataProvider.Factory<T> factory) {
            return addWithFabricOutput(factory::create);
        }

        @Override
        public <T> void addFromCodec(String name, PackOutput.Target target, String directory, Codec<T> codec, Consumer<BiConsumer<Identifier, T>> output) {
            addWithRegistries((out, registries) -> new FabricCodecDataProvider<T>(out, registries, target, directory, codec) {
                @Override
                public String getName() {
                    return name;
                }

                @Override
                protected void configure(BiConsumer<Identifier, T> provider, HolderLookup.Provider registries) {
                    output.accept(provider);
                }
            });
        }

        @Override
        public TagsProvider<Block> blockTags(Consumer<TagProvider.TagConsumer<Block>> tags) {
            return addWithRegistries((out, registries) -> new FabricTagsProvider.BlockTagsProvider(out, registries) {
                @Override
                protected void addTags(HolderLookup.Provider registries) {
                    tags.accept(this::builder);
                }
            });
        }

        @Override
        public TagsProvider<Item> itemTags(Consumer<TagProvider.TagConsumer<Item>> tags) {
            return addWithRegistries((out, registries) -> new FabricTagsProvider.ItemTagsProvider(out, registries) {
                @Override
                protected void addTags(HolderLookup.Provider registries) {
                    tags.accept(this::builder);
                }
            });
        }

        @Override
        public void registries(CompletableFuture<RegistrySetBuilder.PatchedRegistries> registries) {
            addWithFabricOutput(out -> new FabricDynamicRegistryProvider(out, registries.thenApply(RegistrySetBuilder.PatchedRegistries::patches)) {
                @Override
                public String getName() {
                    return "Registries";
                }

                @Override
                protected void configure(HolderLookup.Provider registries, Entries entries) {
                    for (var reg : DynamicRegistries.getBootstrappingRegistries()) {
                        registries.lookupOrThrow(reg.key()).listElements().forEach(x -> register(entries, x));
                    }
                }

                private static <T> void register(Entries entries, Holder.Reference<T> reference) {
                    entries.add(reference.key(), reference.value());
                }
            });
        }

        @Override
        public void addModels(Consumer<BlockModelGenerators> blocks, Consumer<ItemModelGenerators> items) {
            addWithFabricOutput(output -> new FabricModelProvider(output) {
                @Override
                public void generateBlockStateModels(BlockModelGenerators generators) {
                    blocks.accept(generators);
                }

                @Override
                public void generateItemModels(ItemModelGenerators generators) {
                    items.accept(generators);
                }
            });
        }
    }
}
