package net.minecraft.util;

import com.mojang.serialization.Codec;
import com.mojang.serialization.DataResult;
import com.mojang.serialization.DynamicOps;
import com.mojang.serialization.JavaOps;
import com.mojang.serialization.Lifecycle;
import java.util.HashMap;
import java.util.Map;
import java.util.Optional;
import net.minecraft.core.HolderGetter;
import net.minecraft.core.HolderLookup;
import net.minecraft.core.HolderOwner;
import net.minecraft.core.HolderSet;
import net.minecraft.core.Registry;
import net.minecraft.core.Holder.Reference;
import net.minecraft.core.HolderSet.Named;
import net.minecraft.resources.RegistryOps;
import net.minecraft.resources.ResourceKey;
import net.minecraft.tags.TagKey;

public class PlaceholderLookupProvider implements HolderGetter.Provider {
	final HolderLookup.Provider context;
	final PlaceholderLookupProvider.UniversalLookup lookup = new PlaceholderLookupProvider.UniversalLookup();
	final Map<ResourceKey<Object>, Reference<Object>> holders = new HashMap();
	final Map<TagKey<Object>, Named<Object>> holderSets = new HashMap();

	public PlaceholderLookupProvider(HolderLookup.Provider context) {
		this.context = context;
	}

	@Override
	public <T> Optional<? extends HolderGetter<T>> lookup(ResourceKey<? extends Registry<? extends T>> registryKey) {
		return Optional.of(this.lookup.castAsLookup());
	}

	public <V> RegistryOps<V> createSerializationContext(DynamicOps<V> ops) {
		return RegistryOps.create(
			ops,
			new RegistryOps.RegistryInfoLookup() {
				@Override
				public <T> Optional<RegistryOps.RegistryInfo<T>> lookup(ResourceKey<? extends Registry<? extends T>> registryKey) {
					return PlaceholderLookupProvider.this.context
						.lookup(registryKey)
						.map(RegistryOps.RegistryInfo::fromRegistryLookup)
						.or(
							() -> Optional.of(
								new RegistryOps.RegistryInfo(
									PlaceholderLookupProvider.this.lookup.castAsOwner(), PlaceholderLookupProvider.this.lookup.castAsLookup(), Lifecycle.experimental()
								)
							)
						);
				}
			}
		);
	}

	public RegistryContextSwapper createSwapper() {
		return new RegistryContextSwapper() {
			@Override
			public <T> DataResult<T> swapTo(Codec<T> codec, T value, HolderLookup.Provider provider) {
				return codec.encodeStart(PlaceholderLookupProvider.this.createSerializationContext(JavaOps.INSTANCE), value)
					.flatMap(object -> codec.parse(provider.createSerializationContext(JavaOps.INSTANCE), object));
			}
		};
	}

	public boolean hasRegisteredPlaceholders() {
		return !this.holders.isEmpty() || !this.holderSets.isEmpty();
	}

	class UniversalLookup implements HolderGetter<Object>, HolderOwner<Object> {
		@Override
		public Optional<Reference<Object>> get(ResourceKey<Object> resourceKey) {
			return Optional.of(this.getOrCreate(resourceKey));
		}

		@Override
		public Reference<Object> getOrThrow(ResourceKey<Object> resourceKey) {
			return this.getOrCreate(resourceKey);
		}

		private Reference<Object> getOrCreate(ResourceKey<Object> key) {
			return (Reference<Object>)PlaceholderLookupProvider.this.holders.computeIfAbsent(key, resourceKey -> Reference.createStandAlone(this, resourceKey));
		}

		@Override
		public Optional<Named<Object>> get(TagKey<Object> tagKey) {
			return Optional.of(this.getOrCreate(tagKey));
		}

		@Override
		public Named<Object> getOrThrow(TagKey<Object> tagKey) {
			return this.getOrCreate(tagKey);
		}

		private Named<Object> getOrCreate(TagKey<Object> key) {
			return (Named<Object>)PlaceholderLookupProvider.this.holderSets.computeIfAbsent(key, tagKey -> HolderSet.emptyNamed(this, tagKey));
		}

		public <T> HolderGetter<T> castAsLookup() {
			return this;
		}

		public <T> HolderOwner<T> castAsOwner() {
			return this;
		}
	}
}