1 files changed, 18 insertions, 3 deletions
diff --git a/src/main/kotlin/moe/nea/firmament/commands/dsl.kt b/src/main/kotlin/moe/nea/firmament/commands/dsl.kt
index 155e700..90018fa 100644
--- a/src/main/kotlin/moe/nea/firmament/commands/dsl.kt
+++ b/src/main/kotlin/moe/nea/firmament/commands/dsl.kt
@@ -25,6 +25,8 @@ import com.mojang.brigadier.builder.RequiredArgumentBuilder
 import com.mojang.brigadier.context.CommandContext
 import com.mojang.brigadier.suggestion.SuggestionProvider
 import java.lang.reflect.ParameterizedType
+import java.lang.reflect.Type
+import java.lang.reflect.TypeVariable
 import net.fabricmc.fabric.api.client.command.v2.FabricClientCommandSource
 import moe.nea.firmament.util.iterate
 
@@ -43,18 +45,31 @@ fun literal(
 ): LiteralArgumentBuilder<DefaultSource> =
     LiteralArgumentBuilder.literal<DefaultSource>(name).also(block)
 
+
+private fun normalizeGeneric(argument: Type): Class<*> {
+    return if (argument is Class<*>) {
+        argument
+    } else if (argument is TypeVariable<*>) {
+        normalizeGeneric(argument.bounds[0])
+    } else if (argument is ParameterizedType) {
+        normalizeGeneric(argument.rawType)
+    } else {
+        Any::class.java
+    }
+}
+
 data class TypeSafeArg<T : Any>(val name: String, val argument: ArgumentType<T>) {
     val argClass by lazy {
         argument.javaClass
             .iterate<Class<in ArgumentType<T>>> {
                 it.superclass
             }
-            .map {
-                it.genericSuperclass
+            .flatMap {
+                it.genericInterfaces.toList()
             }
             .filterIsInstance<ParameterizedType>()
             .find { it.rawType == ArgumentType::class.java }!!
-            .let { it.actualTypeArguments[0] as Class<*> }
+            .let { normalizeGeneric(it.actualTypeArguments[0]) }
     }
 
     @JvmName("getWithThis")