package net.minecraft.commands.synchronization;

import com.google.common.collect.Maps;
import com.mojang.brigadier.context.CommandContext;
import com.mojang.brigadier.exceptions.CommandSyntaxException;
import com.mojang.brigadier.suggestion.SuggestionProvider;
import com.mojang.brigadier.suggestion.Suggestions;
import com.mojang.brigadier.suggestion.SuggestionsBuilder;
import java.util.Map;
import java.util.concurrent.CompletableFuture;
import net.minecraft.Util;
import net.minecraft.commands.CommandSourceStack;
import net.minecraft.commands.SharedSuggestionProvider;
import net.minecraft.core.registries.BuiltInRegistries;
import net.minecraft.network.chat.Component;
import net.minecraft.resources.ResourceLocation;
import net.minecraft.world.entity.EntityType;

public class SuggestionProviders {
	private static final Map<ResourceLocation, SuggestionProvider<SharedSuggestionProvider>> PROVIDERS_BY_NAME = Maps.<ResourceLocation, SuggestionProvider<SharedSuggestionProvider>>newHashMap();
	private static final ResourceLocation DEFAULT_NAME = ResourceLocation.withDefaultNamespace("ask_server");
	public static final SuggestionProvider<SharedSuggestionProvider> ASK_SERVER = register(
		DEFAULT_NAME, (commandContext, suggestionsBuilder) -> commandContext.getSource().customSuggestion(commandContext)
	);
	public static final SuggestionProvider<CommandSourceStack> AVAILABLE_SOUNDS = register(
		ResourceLocation.withDefaultNamespace("available_sounds"),
		(commandContext, suggestionsBuilder) -> SharedSuggestionProvider.suggestResource(commandContext.getSource().getAvailableSounds(), suggestionsBuilder)
	);
	public static final SuggestionProvider<CommandSourceStack> SUMMONABLE_ENTITIES = register(
		ResourceLocation.withDefaultNamespace("summonable_entities"),
		(commandContext, suggestionsBuilder) -> SharedSuggestionProvider.suggestResource(
			BuiltInRegistries.ENTITY_TYPE.stream().filter(entityType -> entityType.isEnabled(commandContext.getSource().enabledFeatures()) && entityType.canSummon()),
			suggestionsBuilder,
			EntityType::getKey,
			entityType -> Component.translatable(Util.makeDescriptionId("entity", EntityType.getKey(entityType)))
		)
	);

	public static <S extends SharedSuggestionProvider> SuggestionProvider<S> register(ResourceLocation name, SuggestionProvider<SharedSuggestionProvider> provider) {
		if (PROVIDERS_BY_NAME.containsKey(name)) {
			throw new IllegalArgumentException("A command suggestion provider is already registered with the name " + name);
		} else {
			PROVIDERS_BY_NAME.put(name, provider);
			return new SuggestionProviders.Wrapper(name, provider);
		}
	}

	public static SuggestionProvider<SharedSuggestionProvider> getProvider(ResourceLocation name) {
		return (SuggestionProvider<SharedSuggestionProvider>)PROVIDERS_BY_NAME.getOrDefault(name, ASK_SERVER);
	}

	/**
	 * Gets the ID for the given provider. If the provider is not a wrapped one created via {@link #register}, then it returns {@link #ASK_SERVER_ID} instead, as there is no known ID but ASK_SERVER always works.
	 */
	public static ResourceLocation getName(SuggestionProvider<SharedSuggestionProvider> provider) {
		return provider instanceof SuggestionProviders.Wrapper ? ((SuggestionProviders.Wrapper)provider).name : DEFAULT_NAME;
	}

	/**
	 * Checks to make sure that the given suggestion provider is a wrapped one that was created via {@link #register}. If not, returns {@link #ASK_SERVER}. Needed because custom providers don't have a known ID to send to the client, but ASK_SERVER always works.
	 */
	public static SuggestionProvider<SharedSuggestionProvider> safelySwap(SuggestionProvider<SharedSuggestionProvider> provider) {
		return provider instanceof SuggestionProviders.Wrapper ? provider : ASK_SERVER;
	}

	protected static class Wrapper implements SuggestionProvider<SharedSuggestionProvider> {
		private final SuggestionProvider<SharedSuggestionProvider> delegate;
		final ResourceLocation name;

		public Wrapper(ResourceLocation name, SuggestionProvider<SharedSuggestionProvider> delegate) {
			this.delegate = delegate;
			this.name = name;
		}

		@Override
		public CompletableFuture<Suggestions> getSuggestions(CommandContext<SharedSuggestionProvider> commandContext, SuggestionsBuilder suggestionsBuilder) throws CommandSyntaxException {
			return this.delegate.getSuggestions(commandContext, suggestionsBuilder);
		}
	}
}