diff --git a/common/api/src/main/java/com/dfsek/terra/api/command/arguments/RegistryArgument.java b/common/api/src/main/java/com/dfsek/terra/api/command/arguments/RegistryArgument.java index 3e85d17d5..c0f422644 100644 --- a/common/api/src/main/java/com/dfsek/terra/api/command/arguments/RegistryArgument.java +++ b/common/api/src/main/java/com/dfsek/terra/api/command/arguments/RegistryArgument.java @@ -12,34 +12,36 @@ import com.dfsek.terra.api.registry.exception.NoSuchEntryException; import com.dfsek.terra.api.registry.key.RegistryKey; +import com.dfsek.terra.api.structure.Structure; +import com.dfsek.terra.api.util.reflection.TypeKey; + import io.leangen.geantyref.TypeToken; import org.checkerframework.checker.nullness.qual.NonNull; import org.checkerframework.checker.nullness.qual.Nullable; import java.util.List; +import java.util.Optional; import java.util.Queue; import java.util.function.BiFunction; +import java.util.function.Function; import java.util.stream.Collectors; public class RegistryArgument extends CommandArgument { - - - @SuppressWarnings("unchecked") private RegistryArgument( - final boolean required, - final @NonNull String name, - final Registry registry, - final @NonNull String defaultValue, - final @Nullable BiFunction<@NonNull CommandContext, @NonNull String, - @NonNull List<@NonNull String>> suggestionsProvider, - final @NonNull ArgumentDescription description + boolean required, + @NonNull String name, + Function, Registry> registryFunction, + TypeToken typeToken, + @NonNull String defaultValue, + @Nullable BiFunction, String, List> suggestionsProvider, + @NonNull ArgumentDescription description ) { super(required, name, - new RegistryArgumentParser<>(registry), + new RegistryArgumentParser<>(registryFunction), defaultValue, - (TypeToken) TypeToken.get(registry.getType().getType()), + typeToken, suggestionsProvider, description); } @@ -60,13 +62,38 @@ public class RegistryArgument extends CommandArgument { return RegistryArgument.builder(name, registry).asOptionalWithDefault(defaultKey).build(); } + @SuppressWarnings("unchecked") + public static Builder builder(String name, Function, Registry> registryFunction, TypeKey registryType) { + return new Builder<>(name, registryFunction, (TypeToken) TypeToken.get(registryType.getType())); + } + + public static CommandArgument of(String name, Function, Registry> registryFunction, TypeKey registryType) { + return RegistryArgument.builder(name, registryFunction, registryType).build(); + } + + public static CommandArgument optional(String name, Function, Registry> registryFunction, TypeKey registryType) { + return RegistryArgument.builder(name, registryFunction, registryType).asOptional().build(); + } + + public static CommandArgument optional(String name, Function, Registry> registryFunction, TypeKey registryType, String defaultKey) { + return RegistryArgument.builder(name, registryFunction, registryType).asOptionalWithDefault(defaultKey).build(); + } + public static final class Builder extends CommandArgument.Builder { - private final Registry registry; + private final Function, Registry> registryFunction; + private final TypeToken typeToken; @SuppressWarnings("unchecked") private Builder(@NonNull String name, Registry registry) { super((TypeToken) TypeToken.get(registry.getType().getType()), name); - this.registry = registry; + this.registryFunction = commandContext -> registry; + this.typeToken = (TypeToken) TypeToken.get(registry.getType().getType()); + } + + private Builder(@NonNull String name, Function, Registry> registryFunction, TypeToken typeToken) { + super(typeToken, name); + this.typeToken = typeToken; + this.registryFunction = registryFunction; } @Override @@ -74,7 +101,8 @@ public class RegistryArgument extends CommandArgument { return new RegistryArgument<>( isRequired(), getName(), - registry, + registryFunction, + typeToken, getDefaultValue(), getSuggestionsProvider(), getDefaultDescription() @@ -84,22 +112,43 @@ public class RegistryArgument extends CommandArgument { private static final class RegistryArgumentParser implements ArgumentParser { - private final Registry registry; + private final Function, Registry> registryFunction; - private RegistryArgumentParser(Registry registry) { - this.registry = registry; + private RegistryArgumentParser(Function, Registry> registryFunction) { + this.registryFunction = registryFunction; } @Override public @NonNull ArgumentParseResult<@NonNull R> parse(@NonNull CommandContext<@NonNull T> commandContext, @NonNull Queue<@NonNull String> inputQueue) { String input = inputQueue.remove(); - return registry.get(RegistryKey.parse(input)).map(ArgumentParseResult::success).orElse(ArgumentParseResult.failure(new NoSuchEntryException("No such entry: " + input))); + String next = inputQueue.peek(); + if(next != null && next.equals(":")) { + input += inputQueue.remove(); + input += inputQueue.remove(); + } + + Registry registry = registryFunction.apply(commandContext); + + Optional result; + try { + result = registry.get(RegistryKey.parse(input)); + } catch(IllegalArgumentException e) { + try { + result = registry.getByID(input); + } catch(IllegalArgumentException e1) { + return ArgumentParseResult.failure(e1); + } + } + + return result + .map(ArgumentParseResult::success) + .orElse(ArgumentParseResult.failure(new NoSuchEntryException("No such entry: " + input))); } @Override public @NonNull List<@NonNull String> suggestions(@NonNull CommandContext commandContext, @NonNull String input) { - return registry.keys().stream().map(RegistryKey::toString).sorted().collect(Collectors.toList()); + return registryFunction.apply(commandContext).keys().stream().map(RegistryKey::toString).sorted().collect(Collectors.toList()); } } }