package net.minecraft.core;

import com.google.common.collect.ImmutableMap;
import com.google.common.collect.Iterators;
import com.google.common.collect.ImmutableMap.Builder;
import com.mojang.serialization.Lifecycle;
import it.unimi.dsi.fastutil.objects.ObjectArrayList;
import it.unimi.dsi.fastutil.objects.ObjectList;
import it.unimi.dsi.fastutil.objects.Reference2IntMap;
import it.unimi.dsi.fastutil.objects.Reference2IntOpenHashMap;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.IdentityHashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Optional;
import java.util.Set;
import java.util.Map.Entry;
import java.util.function.BiConsumer;
import java.util.stream.Stream;
import net.minecraft.Util;
import net.minecraft.core.Holder.Reference;
import net.minecraft.core.HolderLookup.RegistryLookup.Delegate;
import net.minecraft.core.HolderSet.Named;
import net.minecraft.core.MappedRegistry.TagSet.1;
import net.minecraft.core.MappedRegistry.TagSet.2;
import net.minecraft.core.Registry.PendingTags;
import net.minecraft.resources.ResourceKey;
import net.minecraft.resources.ResourceLocation;
import net.minecraft.tags.TagKey;
import net.minecraft.tags.TagLoader.LoadResult;
import net.minecraft.util.RandomSource;
import org.jetbrains.annotations.Nullable;

public class MappedRegistry<T> implements WritableRegistry<T> {
	private final ResourceKey<? extends Registry<T>> key;
	private final ObjectList<Reference<T>> byId = new ObjectArrayList<>(256);
	private final Reference2IntMap<T> toId = Util.make(
		new Reference2IntOpenHashMap<>(), reference2IntOpenHashMap -> reference2IntOpenHashMap.defaultReturnValue(-1)
	);
	private final Map<ResourceLocation, Reference<T>> byLocation = new HashMap();
	private final Map<ResourceKey<T>, Reference<T>> byKey = new HashMap();
	private final Map<T, Reference<T>> byValue = new IdentityHashMap();
	private final Map<ResourceKey<T>, RegistrationInfo> registrationInfos = new IdentityHashMap();
	private Lifecycle registryLifecycle;
	private final Map<TagKey<T>, Named<T>> frozenTags = new IdentityHashMap();
	MappedRegistry.TagSet<T> allTags = MappedRegistry.TagSet.unbound();
	private boolean frozen;
	@Nullable
	private Map<T, Reference<T>> unregisteredIntrusiveHolders;

	@Override
	public Stream<Named<T>> listTags() {
		return this.getTags();
	}

	public MappedRegistry(ResourceKey<? extends Registry<T>> key, Lifecycle registryLifecycle) {
		this(key, registryLifecycle, false);
	}

	public MappedRegistry(ResourceKey<? extends Registry<T>> key, Lifecycle registryLifecycle, boolean hasIntrusiveHolders) {
		this.key = key;
		this.registryLifecycle = registryLifecycle;
		if (hasIntrusiveHolders) {
			this.unregisteredIntrusiveHolders = new IdentityHashMap();
		}
	}

	@Override
	public ResourceKey<? extends Registry<T>> key() {
		return this.key;
	}

	public String toString() {
		return "Registry[" + this.key + " (" + this.registryLifecycle + ")]";
	}

	private void validateWrite() {
		if (this.frozen) {
			throw new IllegalStateException("Registry is already frozen");
		}
	}

	private void validateWrite(ResourceKey<T> key) {
		if (this.frozen) {
			throw new IllegalStateException("Registry is already frozen (trying to add key " + key + ")");
		}
	}

	@Override
	public Reference<T> register(ResourceKey<T> key, T value, RegistrationInfo registrationInfo) {
		this.validateWrite(key);
		Objects.requireNonNull(key);
		Objects.requireNonNull(value);
		if (this.byLocation.containsKey(key.location())) {
			throw (IllegalStateException)Util.pauseInIde((T)(new IllegalStateException("Adding duplicate key '" + key + "' to registry")));
		} else if (this.byValue.containsKey(value)) {
			throw (IllegalStateException)Util.pauseInIde((T)(new IllegalStateException("Adding duplicate value '" + value + "' to registry")));
		} else {
			Reference<T> reference;
			if (this.unregisteredIntrusiveHolders != null) {
				reference = (Reference<T>)this.unregisteredIntrusiveHolders.remove(value);
				if (reference == null) {
					throw new AssertionError("Missing intrusive holder for " + key + ":" + value);
				}

				reference.bindKey(key);
			} else {
				reference = (Reference<T>)this.byKey.computeIfAbsent(key, resourceKey -> Reference.createStandAlone(this, resourceKey));
			}

			this.byKey.put(key, reference);
			this.byLocation.put(key.location(), reference);
			this.byValue.put(value, reference);
			int i = this.byId.size();
			this.byId.add(reference);
			this.toId.put(value, i);
			this.registrationInfos.put(key, registrationInfo);
			this.registryLifecycle = this.registryLifecycle.add(registrationInfo.lifecycle());
			return reference;
		}
	}

	@Nullable
	@Override
	public ResourceLocation getKey(T value) {
		Reference<T> reference = (Reference<T>)this.byValue.get(value);
		return reference != null ? reference.key().location() : null;
	}

	@Override
	public Optional<ResourceKey<T>> getResourceKey(T value) {
		return Optional.ofNullable((Reference)this.byValue.get(value)).map(Reference::key);
	}

	@Override
	public int getId(@Nullable T value) {
		return this.toId.getInt(value);
	}

	@Nullable
	@Override
	public T getValue(@Nullable ResourceKey<T> key) {
		return getValueFromNullable((Reference<T>)this.byKey.get(key));
	}

	@Nullable
	@Override
	public T byId(int id) {
		return (T)(id >= 0 && id < this.byId.size() ? ((Reference)this.byId.get(id)).value() : null);
	}

	@Override
	public Optional<Reference<T>> get(int index) {
		return index >= 0 && index < this.byId.size() ? Optional.ofNullable((Reference)this.byId.get(index)) : Optional.empty();
	}

	@Override
	public Optional<Reference<T>> get(ResourceLocation key) {
		return Optional.ofNullable((Reference)this.byLocation.get(key));
	}

	@Override
	public Optional<Reference<T>> get(ResourceKey<T> resourceKey) {
		return Optional.ofNullable((Reference)this.byKey.get(resourceKey));
	}

	@Override
	public Optional<Reference<T>> getAny() {
		return this.byId.isEmpty() ? Optional.empty() : Optional.of((Reference)this.byId.getFirst());
	}

	@Override
	public Holder<T> wrapAsHolder(T value) {
		Reference<T> reference = (Reference<T>)this.byValue.get(value);
		return (Holder<T>)(reference != null ? reference : Holder.direct(value));
	}

	Reference<T> getOrCreateHolderOrThrow(ResourceKey<T> key) {
		return (Reference<T>)this.byKey.computeIfAbsent(key, resourceKey -> {
			if (this.unregisteredIntrusiveHolders != null) {
				throw new IllegalStateException("This registry can't create new holders without value");
			} else {
				this.validateWrite(resourceKey);
				return Reference.createStandAlone(this, resourceKey);
			}
		});
	}

	@Override
	public int size() {
		return this.byKey.size();
	}

	@Override
	public Optional<RegistrationInfo> registrationInfo(ResourceKey<T> key) {
		return Optional.ofNullable((RegistrationInfo)this.registrationInfos.get(key));
	}

	@Override
	public Lifecycle registryLifecycle() {
		return this.registryLifecycle;
	}

	public Iterator<T> iterator() {
		return Iterators.transform(this.byId.iterator(), Holder::value);
	}

	@Nullable
	@Override
	public T getValue(@Nullable ResourceLocation key) {
		Reference<T> reference = (Reference<T>)this.byLocation.get(key);
		return getValueFromNullable(reference);
	}

	@Nullable
	private static <T> T getValueFromNullable(@Nullable Reference<T> holder) {
		return holder != null ? holder.value() : null;
	}

	@Override
	public Set<ResourceLocation> keySet() {
		return Collections.unmodifiableSet(this.byLocation.keySet());
	}

	@Override
	public Set<ResourceKey<T>> registryKeySet() {
		return Collections.unmodifiableSet(this.byKey.keySet());
	}

	@Override
	public Set<Entry<ResourceKey<T>, T>> entrySet() {
		return Collections.unmodifiableSet(Util.mapValuesLazy(this.byKey, Holder::value).entrySet());
	}

	@Override
	public Stream<Reference<T>> listElements() {
		return this.byId.stream();
	}

	@Override
	public Stream<Named<T>> getTags() {
		return this.allTags.getTags();
	}

	Named<T> getOrCreateTagForRegistration(TagKey<T> key) {
		return (Named<T>)this.frozenTags.computeIfAbsent(key, this::createTag);
	}

	private Named<T> createTag(TagKey<T> key) {
		return new Named<>(this, key);
	}

	@Override
	public boolean isEmpty() {
		return this.byKey.isEmpty();
	}

	@Override
	public Optional<Reference<T>> getRandom(RandomSource random) {
		return Util.getRandomSafe(this.byId, random);
	}

	@Override
	public boolean containsKey(ResourceLocation name) {
		return this.byLocation.containsKey(name);
	}

	@Override
	public boolean containsKey(ResourceKey<T> key) {
		return this.byKey.containsKey(key);
	}

	@Override
	public Registry<T> freeze() {
		if (this.frozen) {
			return this;
		} else {
			this.frozen = true;
			this.byValue.forEach((object, reference) -> reference.bindValue(object));
			List<ResourceLocation> list = this.byKey
				.entrySet()
				.stream()
				.filter(entry -> !((Reference)entry.getValue()).isBound())
				.map(entry -> ((ResourceKey)entry.getKey()).location())
				.sorted()
				.toList();
			if (!list.isEmpty()) {
				throw new IllegalStateException("Unbound values in registry " + this.key() + ": " + list);
			} else {
				if (this.unregisteredIntrusiveHolders != null) {
					if (!this.unregisteredIntrusiveHolders.isEmpty()) {
						throw new IllegalStateException("Some intrusive holders were not registered: " + this.unregisteredIntrusiveHolders.values());
					}

					this.unregisteredIntrusiveHolders = null;
				}

				if (this.allTags.isBound()) {
					throw new IllegalStateException("Tags already present before freezing");
				} else {
					List<ResourceLocation> list2 = this.frozenTags
						.entrySet()
						.stream()
						.filter(entry -> !((Named)entry.getValue()).isBound())
						.map(entry -> ((TagKey)entry.getKey()).location())
						.sorted()
						.toList();
					if (!list2.isEmpty()) {
						throw new IllegalStateException("Unbound tags in registry " + this.key() + ": " + list2);
					} else {
						this.allTags = MappedRegistry.TagSet.fromMap(this.frozenTags);
						this.refreshTagsInHolders();
						return this;
					}
				}
			}
		}
	}

	@Override
	public Reference<T> createIntrusiveHolder(T value) {
		if (this.unregisteredIntrusiveHolders == null) {
			throw new IllegalStateException("This registry can't create intrusive holders");
		} else {
			this.validateWrite();
			return (Reference<T>)this.unregisteredIntrusiveHolders.computeIfAbsent(value, object -> Reference.createIntrusive(this, (T)object));
		}
	}

	@Override
	public Optional<Named<T>> get(TagKey<T> tagKey) {
		return this.allTags.get(tagKey);
	}

	private Reference<T> validateAndUnwrapTagElement(TagKey<T> key, Holder<T> value) {
		if (!value.canSerializeIn(this)) {
			throw new IllegalStateException("Can't create named set " + key + " containing value " + value + " from outside registry " + this);
		} else if (value instanceof Reference<T> reference) {
			return reference;
		} else {
			throw new IllegalStateException("Found direct holder " + value + " value in tag " + key);
		}
	}

	@Override
	public void bindTag(TagKey<T> tag, List<Holder<T>> values) {
		this.validateWrite();
		this.getOrCreateTagForRegistration(tag).bind(values);
	}

	void refreshTagsInHolders() {
		Map<Reference<T>, List<TagKey<T>>> map = new IdentityHashMap();
		this.byKey.values().forEach(reference -> map.put(reference, new ArrayList()));
		this.allTags.forEach((tagKey, named) -> {
			for (Holder<T> holder : named) {
				Reference<T> reference = this.validateAndUnwrapTagElement(tagKey, holder);
				((List)map.get(reference)).add(tagKey);
			}
		});
		map.forEach(Reference::bindTags);
	}

	public void bindAllTagsToEmpty() {
		this.validateWrite();
		this.frozenTags.values().forEach(named -> named.bind(List.of()));
	}

	@Override
	public HolderGetter<T> createRegistrationLookup() {
		this.validateWrite();
		return new HolderGetter<T>() {
			@Override
			public Optional<Reference<T>> get(ResourceKey<T> resourceKey) {
				return Optional.of(this.getOrThrow(resourceKey));
			}

			@Override
			public Reference<T> getOrThrow(ResourceKey<T> resourceKey) {
				return MappedRegistry.this.getOrCreateHolderOrThrow(resourceKey);
			}

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

			@Override
			public Named<T> getOrThrow(TagKey<T> tagKey) {
				return MappedRegistry.this.getOrCreateTagForRegistration(tagKey);
			}
		};
	}

	@Override
	public PendingTags<T> prepareTagReload(LoadResult<T> loadResult) {
		if (!this.frozen) {
			throw new IllegalStateException("Invalid method used for tag loading");
		} else {
			Builder<TagKey<T>, Named<T>> builder = ImmutableMap.builder();
			final Map<TagKey<T>, List<Holder<T>>> map = new HashMap();
			loadResult.tags().forEach((tagKey, list) -> {
				Named<T> named = (Named<T>)this.frozenTags.get(tagKey);
				if (named == null) {
					named = this.createTag(tagKey);
				}

				builder.put(tagKey, named);
				map.put(tagKey, List.copyOf(list));
			});
			final ImmutableMap<TagKey<T>, Named<T>> immutableMap = builder.build();
			final HolderLookup.RegistryLookup<T> registryLookup = new Delegate<T>() {
				public HolderLookup.RegistryLookup<T> parent() {
					return MappedRegistry.this;
				}

				public Optional<Named<T>> get(TagKey<T> tagKey) {
					return Optional.ofNullable(immutableMap.get(tagKey));
				}

				public Stream<Named<T>> listTags() {
					return immutableMap.values().stream();
				}
			};
			return new PendingTags<T>() {
				@Override
				public ResourceKey<? extends Registry<? extends T>> key() {
					return MappedRegistry.this.key();
				}

				@Override
				public int size() {
					return map.size();
				}

				@Override
				public HolderLookup.RegistryLookup<T> lookup() {
					return registryLookup;
				}

				@Override
				public void apply() {
					immutableMap.forEach((tagKey, named) -> {
						List<Holder<T>> list = (List<Holder<T>>)map.getOrDefault(tagKey, List.of());
						named.bind(list);
					});
					MappedRegistry.this.allTags = MappedRegistry.TagSet.fromMap(immutableMap);
					MappedRegistry.this.refreshTagsInHolders();
				}
			};
		}
	}

	interface TagSet<T> {
		static <T> MappedRegistry.TagSet<T> unbound() {
			return new 1();
		}

		static <T> MappedRegistry.TagSet<T> fromMap(Map<TagKey<T>, Named<T>> map) {
			return new 2(map);
		}

		boolean isBound();

		Optional<Named<T>> get(TagKey<T> key);

		void forEach(BiConsumer<? super TagKey<T>, ? super Named<T>> action);

		Stream<Named<T>> getTags();
	}
}