add api for RegistryArgument to get registry at runtime

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<T, R> extends CommandArgument<T, R> {
-    
-    
-    @SuppressWarnings("unchecked")
     private RegistryArgument(
-            final boolean required,
+            boolean required,
-            final @NonNull String name,
+            @NonNull String name,
-            final Registry<R> registry,
+            Function<CommandContext<T>, Registry<R>> registryFunction,
-            final @NonNull String defaultValue,
+            TypeToken<R> typeToken,
-            final @Nullable BiFunction<@NonNull CommandContext<T>, @NonNull String,
+            @NonNull String defaultValue,
-                    @NonNull List<@NonNull String>> suggestionsProvider,
+            @Nullable BiFunction<CommandContext<T>, String, List<String>> suggestionsProvider,
-            final @NonNull ArgumentDescription description
+            @NonNull ArgumentDescription description
                             ) {
         super(required,
               name,
-              new RegistryArgumentParser<>(registry),
+              new RegistryArgumentParser<>(registryFunction),
               defaultValue,
-              (TypeToken<R>) TypeToken.get(registry.getType().getType()),
+              typeToken,
               suggestionsProvider,
               description);
     }
@@ -60,13 +62,38 @@ public class RegistryArgument<T, R> extends CommandArgument<T, R> {
         return RegistryArgument.<T, R>builder(name, registry).asOptionalWithDefault(defaultKey).build();
     }
     
+    @SuppressWarnings("unchecked")
+    public static <T, R> Builder<T, R> builder(String name, Function<CommandContext<T>, Registry<R>> registryFunction, TypeKey<R> registryType) {
+        return new Builder<>(name, registryFunction, (TypeToken<R>) TypeToken.get(registryType.getType()));
+    }
+    
+    public static <T, R> CommandArgument<T, R> of(String name, Function<CommandContext<T>, Registry<R>> registryFunction, TypeKey<R> registryType) {
+        return RegistryArgument.<T, R>builder(name, registryFunction, registryType).build();
+    }
+    
+    public static <T, R> CommandArgument<T, R> optional(String name, Function<CommandContext<T>, Registry<R>> registryFunction, TypeKey<R> registryType) {
+        return RegistryArgument.builder(name, registryFunction, registryType).asOptional().build();
+    }
+    
+    public static <T, R> CommandArgument<T, R> optional(String name, Function<CommandContext<T>, Registry<R>> registryFunction, TypeKey<R> registryType, String defaultKey) {
+        return RegistryArgument.builder(name, registryFunction, registryType).asOptionalWithDefault(defaultKey).build();
+    }
+    
     public static final class Builder<T, R> extends CommandArgument.Builder<T, R> {
-        private final Registry<R> registry;
+        private final Function<CommandContext<T>, Registry<R>> registryFunction;
+        private final TypeToken<R> typeToken;
         
         @SuppressWarnings("unchecked")
         private Builder(@NonNull String name, Registry<R> registry) {
             super((TypeToken<R>) TypeToken.get(registry.getType().getType()), name);
-            this.registry = registry;
+            this.registryFunction = commandContext -> registry;
+            this.typeToken = (TypeToken<R>) TypeToken.get(registry.getType().getType());
+        }
+    
+        private Builder(@NonNull String name, Function<CommandContext<T>, Registry<R>> registryFunction, TypeToken<R> typeToken) {
+            super(typeToken, name);
+            this.typeToken = typeToken;
+            this.registryFunction = registryFunction;
         }
         
         @Override
@@ -74,7 +101,8 @@ public class RegistryArgument<T, R> extends CommandArgument<T, R> {
             return new RegistryArgument<>(
                     isRequired(),
                     getName(),
-                    registry,
+                    registryFunction,
+                    typeToken,
                     getDefaultValue(),
                     getSuggestionsProvider(),
                     getDefaultDescription()
@@ -84,22 +112,43 @@ public class RegistryArgument<T, R> extends CommandArgument<T, R> {
     
     
     private static final class RegistryArgumentParser<T, R> implements ArgumentParser<T, R> {
-        private final Registry<R> registry;
+        private final Function<CommandContext<T>, Registry<R>> registryFunction;
         
-        private RegistryArgumentParser(Registry<R> registry) {
+        private RegistryArgumentParser(Function<CommandContext<T>, Registry<R>> registryFunction) {
-            this.registry = registry;
+            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<R> registry = registryFunction.apply(commandContext);
+            
+            Optional<R> 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<T> 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());
         }
     }
 }