package net.minecraft.core;

import com.mojang.serialization.DynamicOps;
import io.netty.buffer.ByteBuf;
import java.util.ArrayList;
import java.util.List;
import java.util.Optional;
import java.util.Set;
import java.util.function.BiConsumer;
import java.util.stream.Collectors;
import java.util.stream.Stream;
import net.minecraft.nbt.Tag;
import net.minecraft.network.codec.ByteBufCodecs;
import net.minecraft.network.codec.StreamCodec;
import net.minecraft.resources.RegistryDataLoader;
import net.minecraft.resources.ResourceKey;
import net.minecraft.resources.ResourceLocation;
import net.minecraft.server.RegistryLayer;
import net.minecraft.server.packs.repository.KnownPack;

public class RegistrySynchronization {
	private static final Set<ResourceKey<? extends Registry<?>>> NETWORKABLE_REGISTRIES = (Set<ResourceKey<? extends Registry<?>>>)RegistryDataLoader.SYNCHRONIZED_REGISTRIES
		.stream()
		.map(RegistryDataLoader.RegistryData::key)
		.collect(Collectors.toUnmodifiableSet());

	public static void packRegistries(
		DynamicOps<Tag> ops,
		RegistryAccess registryAccess,
		Set<KnownPack> packs,
		BiConsumer<ResourceKey<? extends Registry<?>>, List<RegistrySynchronization.PackedRegistryEntry>> packetSender
	) {
		RegistryDataLoader.SYNCHRONIZED_REGISTRIES.forEach(registryData -> packRegistry(ops, registryData, registryAccess, packs, packetSender));
	}

	private static <T> void packRegistry(
		DynamicOps<Tag> ops,
		RegistryDataLoader.RegistryData<T> registryData,
		RegistryAccess registryAccess,
		Set<KnownPack> packs,
		BiConsumer<ResourceKey<? extends Registry<?>>, List<RegistrySynchronization.PackedRegistryEntry>> packetSender
	) {
		registryAccess.lookup(registryData.key())
			.ifPresent(
				registry -> {
					List<RegistrySynchronization.PackedRegistryEntry> list = new ArrayList(registry.size());
					registry.listElements()
						.forEach(
							reference -> {
								boolean bl = registry.registrationInfo(reference.key()).flatMap(RegistrationInfo::knownPackInfo).filter(packs::contains).isPresent();
								Optional<Tag> optional;
								if (bl) {
									optional = Optional.empty();
								} else {
									Tag tag = registryData.elementCodec()
										.encodeStart(ops, (T)reference.value())
										.getOrThrow(string -> new IllegalArgumentException("Failed to serialize " + reference.key() + ": " + string));
									optional = Optional.of(tag);
								}

								list.add(new RegistrySynchronization.PackedRegistryEntry(reference.key().location(), optional));
							}
						);
					packetSender.accept(registry.key(), list);
				}
			);
	}

	private static Stream<RegistryAccess.RegistryEntry<?>> ownedNetworkableRegistries(RegistryAccess registryAccess) {
		return registryAccess.registries().filter(registryEntry -> isNetworkable(registryEntry.key()));
	}

	public static Stream<RegistryAccess.RegistryEntry<?>> networkedRegistries(LayeredRegistryAccess<RegistryLayer> registryAccess) {
		return ownedNetworkableRegistries(registryAccess.getAccessFrom(RegistryLayer.WORLDGEN));
	}

	public static Stream<RegistryAccess.RegistryEntry<?>> networkSafeRegistries(LayeredRegistryAccess<RegistryLayer> registryAccess) {
		Stream<RegistryAccess.RegistryEntry<?>> stream = registryAccess.getLayer(RegistryLayer.STATIC).registries();
		Stream<RegistryAccess.RegistryEntry<?>> stream2 = networkedRegistries(registryAccess);
		return Stream.concat(stream2, stream);
	}

	public static boolean isNetworkable(ResourceKey<? extends Registry<?>> registryKey) {
		return NETWORKABLE_REGISTRIES.contains(registryKey);
	}

	public record PackedRegistryEntry(ResourceLocation id, Optional<Tag> data) {
		public static final StreamCodec<ByteBuf, RegistrySynchronization.PackedRegistryEntry> STREAM_CODEC = StreamCodec.composite(
			ResourceLocation.STREAM_CODEC,
			RegistrySynchronization.PackedRegistryEntry::id,
			ByteBufCodecs.TAG.apply(ByteBufCodecs::optional),
			RegistrySynchronization.PackedRegistryEntry::data,
			RegistrySynchronization.PackedRegistryEntry::new
		);
	}
}