package net.minecraft.commands.synchronization;

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.HashMap;
import java.util.Map;
import java.util.concurrent.CompletableFuture;
import net.minecraft.commands.SharedSuggestionProvider;
import net.minecraft.core.registries.BuiltInRegistries;
import net.minecraft.resources.ResourceLocation;
import net.minecraft.world.entity.EntityType;

public class SuggestionProviders {
    private static final Map<ResourceLocation, SuggestionProvider<SharedSuggestionProvider>> PROVIDERS_BY_NAME = new HashMap<>();
    private static final ResourceLocation ID_ASK_SERVER = ResourceLocation.withDefaultNamespace("ask_server");
    public static final SuggestionProvider<SharedSuggestionProvider> ASK_SERVER = register(
        ID_ASK_SERVER, (context, builder) -> context.getSource().customSuggestion(context)
    );
    public static final SuggestionProvider<SharedSuggestionProvider> AVAILABLE_SOUNDS = register(
        ResourceLocation.withDefaultNamespace("available_sounds"),
        (context, builder) -> SharedSuggestionProvider.suggestResource(context.getSource().getAvailableSounds(), builder)
    );
    public static final SuggestionProvider<SharedSuggestionProvider> 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::getDescription
        )
    );

    public static <S extends SharedSuggestionProvider> SuggestionProvider<S> register(
        ResourceLocation name, SuggestionProvider<SharedSuggestionProvider> provider
    ) {
        SuggestionProvider<SharedSuggestionProvider> suggestionProvider = PROVIDERS_BY_NAME.putIfAbsent(name, provider);
        if (suggestionProvider != null) {
            throw new IllegalArgumentException("A command suggestion provider is already registered with the name '" + name + "'");
        } else {
            return (SuggestionProvider<S>) new SuggestionProviders.RegisteredSuggestion(name, provider);
        }
    }

    public static <S extends SharedSuggestionProvider> SuggestionProvider<S> cast(SuggestionProvider<SharedSuggestionProvider> provider) {
        return (SuggestionProvider<S>)provider;
    }

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

    public static ResourceLocation getName(SuggestionProvider<?> provider) {
        return provider instanceof SuggestionProviders.RegisteredSuggestion registeredSuggestion ? registeredSuggestion.name : ID_ASK_SERVER;
    }

    record RegisteredSuggestion(ResourceLocation name, SuggestionProvider<SharedSuggestionProvider> delegate)
        implements SuggestionProvider<SharedSuggestionProvider> {
        @Override
        public CompletableFuture<Suggestions> getSuggestions(CommandContext<SharedSuggestionProvider> context, SuggestionsBuilder builder) throws CommandSyntaxException {
            return this.delegate.getSuggestions(context, builder);
        }
    }
}
