feat: nested sealed hierarchies for oneOf/anyOf subtypes

core/src/main/kotlin/com/avsystem/justworks/core/gen/ClientGenerator.kt

@@ -31,7 +31,11 @@ private const val API_SUFFIX = "Api"
  * that extends `ApiClientBase` with suspend functions for every endpoint in that tag group.
  */
 @OptIn(ExperimentalKotlinPoetApi::class)
-class ClientGenerator(private val apiPackage: String, private val modelPackage: String) {
+class ClientGenerator(
+    private val apiPackage: String,
+    private val modelPackage: String,
+    private val classNameLookup: Map<String, ClassName> = emptyMap(),
+) {
     fun generate(spec: ApiSpec, hasPolymorphicTypes: Boolean = false): List<FileSpec> {
         val grouped = spec.endpoints.groupBy { it.tags.firstOrNull() ?: DEFAULT_TAG }
         return grouped.map { (tag, endpoints) -> generateClientFile(tag, endpoints, hasPolymorphicTypes) }
@@ -95,7 +99,7 @@ class ClientGenerator(private val apiPackage: String, private val modelPackage:
         val params = endpoint.parameters.groupBy { it.location }
 
         val pathParams = params[ParameterLocation.PATH].orEmpty().map { param ->
-            ParameterSpec(param.name.toCamelCase(), TypeMapping.toTypeName(param.schema, modelPackage))
+            ParameterSpec(param.name.toCamelCase(), TypeMapping.toTypeName(param.schema, modelPackage, classNameLookup))
         }
 
         val queryParams = params[ParameterLocation.QUERY].orEmpty().map { param ->
@@ -124,7 +128,7 @@ class ClientGenerator(private val apiPackage: String, private val modelPackage:
         name: String,
         required: Boolean,
     ): ParameterSpec {
-        val baseType = TypeMapping.toTypeName(typeRef, modelPackage)
+        val baseType = TypeMapping.toTypeName(typeRef, modelPackage, classNameLookup)
 
         val builder = ParameterSpec.builder(name.toCamelCase(), baseType.copy(nullable = !required))
         if (!required) builder.defaultValue("null")
@@ -202,7 +206,7 @@ class ClientGenerator(private val apiPackage: String, private val modelPackage:
         .asSequence()
         .filter { it.key.startsWith("2") }
         .firstNotNullOfOrNull { it.value.schema }
-        ?.let { successResponse -> TypeMapping.toTypeName(successResponse, modelPackage) }
+        ?.let { successResponse -> TypeMapping.toTypeName(successResponse, modelPackage, classNameLookup) }
         ?: UNIT
 
     /**

core/src/main/kotlin/com/avsystem/justworks/core/gen/CodeGenerator.kt

@@ -20,8 +20,10 @@ object CodeGenerator {
         modelFiles.forEach { it.writeTo(outputDir) }
 
         val hasPolymorphicTypes = modelFiles.any { it.name == SerializersModuleGenerator.FILE_NAME }
+        val classNameLookup = ModelGenerator.buildClassNameLookup(spec, modelPackage)
 
-        val clientFiles = ClientGenerator(apiPackage, modelPackage).generate(spec, hasPolymorphicTypes)
+        val clientFiles = ClientGenerator(apiPackage, modelPackage, classNameLookup)
+            .generate(spec, hasPolymorphicTypes)
         clientFiles.forEach { it.writeTo(outputDir) }
 
         return Result(modelFiles.size, clientFiles.size)

core/src/main/kotlin/com/avsystem/justworks/core/gen/ModelGenerator.kt

@@ -31,21 +31,40 @@ import kotlin.time.Instant
  * and one file per [EnumModel] (enum class), all annotated with kotlinx.serialization annotations.
  */
 class ModelGenerator(private val modelPackage: String) {
-    fun generate(spec: ApiSpec): List<FileSpec> = context(
-        buildHierarchyInfo(spec.schemas),
-        InlineSchemaDeduplicator(spec.schemas.map { it.name }.toSet()),
-    ) {
-        val schemaFiles = spec.schemas.flatMap { generateSchemaFiles(it) }
+    fun generate(spec: ApiSpec): List<FileSpec> {
+        val hierarchyInfo = buildHierarchyInfo(spec.schemas)
+        return context(
+            hierarchyInfo,
+            InlineSchemaDeduplicator(spec.schemas.map { it.name }.toSet()),
+        ) {
+            val variantNames = hierarchyInfo.sealedHierarchies
+                .filterKeys { it !in hierarchyInfo.anyOfWithoutDiscriminator }
+                .values
+                .flatten()
+                .toSet()
+
+            val schemaFiles = spec.schemas
+                .filter { it.name !in variantNames }
+                .flatMap { generateSchemaFiles(it) }
+
+            val inlineSchemaFiles = collectAllInlineSchemas(spec).map {
+                if (it.isNested) generateNestedInlineClass(it) else generateDataClass(it)
+            }
 
-        val inlineSchemaFiles = collectAllInlineSchemas(spec).map {
-            if (it.isNested) generateNestedInlineClass(it) else generateDataClass(it)
-        }
+            val enumFiles = spec.enums.map(::generateEnumClass)
 
-        val enumFiles = spec.enums.map(::generateEnumClass)
+            val serializersModuleFile = SerializersModuleGenerator(modelPackage).generate()
 
-        val serializersModuleFile = SerializersModuleGenerator(modelPackage).generate()
+            schemaFiles + inlineSchemaFiles + enumFiles + listOfNotNull(serializersModuleFile)
+        }
+    }
 
-        schemaFiles + inlineSchemaFiles + enumFiles + listOfNotNull(serializersModuleFile)
+    companion object {
+        fun buildClassNameLookup(spec: ApiSpec, modelPackage: String): Map<String, ClassName> {
+            val generator = ModelGenerator(modelPackage)
+            val hierarchy = generator.buildHierarchyInfo(spec.schemas)
+            return generator.buildClassNameLookup(hierarchy)
+        }
     }
 
     data class HierarchyInfo(
@@ -88,6 +107,19 @@ class ModelGenerator(private val modelPackage: String) {
         return HierarchyInfo(sealedHierarchies, variantParents, anyOfWithoutDiscriminator, schemas)
     }
 
+    private fun buildClassNameLookup(hierarchy: HierarchyInfo): Map<String, ClassName> {
+        val lookup = mutableMapOf<String, ClassName>()
+        for ((parentName, variants) in hierarchy.sealedHierarchies) {
+            if (parentName in hierarchy.anyOfWithoutDiscriminator) continue
+            val parentClass = ClassName(modelPackage, parentName)
+            lookup[parentName] = parentClass
+            for (variantName in variants) {
+                lookup[variantName] = parentClass.nestedClass(variantName)
+            }
+        }
+        return lookup
+    }
+
     context(deduplicator: InlineSchemaDeduplicator)
     private fun collectAllInlineSchemas(spec: ApiSpec): List<SchemaModel> {
         val endpointRefs = spec.endpoints.flatMap { endpoint ->
@@ -122,7 +154,7 @@ class ModelGenerator(private val modelPackage: String) {
             if (schema.name in hierarchy.anyOfWithoutDiscriminator) {
                 listOf(generateSealedInterface(schema), generatePolymorphicSerializer(schema))
             } else {
-                listOf(generateSealedInterface(schema))
+                listOf(generateSealedHierarchy(schema))
             }
         }
 
@@ -185,6 +217,107 @@ class ModelGenerator(private val modelPackage: String) {
         return fileBuilder.build()
     }
 
+    /**
+     * Generates a sealed class with nested data class subtypes for oneOf and anyOf-with-discriminator schemas.
+     * Produces a single FileSpec containing the entire hierarchy.
+     */
+    context(hierarchy: HierarchyInfo)
+    private fun generateSealedHierarchy(schema: SchemaModel): FileSpec {
+        val className = ClassName(modelPackage, schema.name)
+        val variants = hierarchy.sealedHierarchies[schema.name].orEmpty()
+        val schemasById = hierarchy.schemas.associateBy { it.name }
+
+        val parentBuilder = TypeSpec.classBuilder(className).addModifiers(KModifier.SEALED)
+        parentBuilder.addAnnotation(SERIALIZABLE)
+
+        if (schema.discriminator != null) {
+            parentBuilder.addAnnotation(
+                AnnotationSpec
+                    .builder(JSON_CLASS_DISCRIMINATOR)
+                    .addMember("%S", schema.discriminator.propertyName)
+                    .build(),
+            )
+        }
+
+        if (schema.description != null) {
+            parentBuilder.addKdoc("%L", schema.description)
+        }
+
+        for (variantName in variants) {
+            val variantSchema = schemasById[variantName] ?: continue
+            val serialName = resolveSerialName(schema, variantName)
+            val nestedType = buildNestedDataClass(variantSchema, className, serialName)
+            parentBuilder.addType(nestedType)
+        }
+
+        val fileBuilder = FileSpec.builder(className).addType(parentBuilder.build())
+
+        if (schema.discriminator != null) {
+            fileBuilder.addAnnotation(
+                AnnotationSpec
+                    .builder(OPT_IN)
+                    .addMember("%T::class", EXPERIMENTAL_SERIALIZATION_API)
+                    .build(),
+            )
+        }
+
+        return fileBuilder.build()
+    }
+
+    /**
+     * Builds a nested data class TypeSpec for a variant within a sealed hierarchy.
+     */
+    private fun buildNestedDataClass(
+        variantSchema: SchemaModel,
+        parentClassName: ClassName,
+        serialName: String,
+    ): TypeSpec {
+        val sortedProps = variantSchema.properties.sortedBy { prop ->
+            when {
+                prop.name in variantSchema.requiredProperties && prop.defaultValue == null -> 1
+                prop.defaultValue != null -> 2
+                else -> 3
+            }
+        }
+
+        val constructorBuilder = FunSpec.constructorBuilder()
+        val propertySpecs = sortedProps.map { prop ->
+            val type = TypeMapping.toTypeName(prop.type, modelPackage).copy(nullable = prop.nullable)
+            val kotlinName = prop.name.toCamelCase()
+
+            val paramBuilder = ParameterSpec.builder(kotlinName, type)
+
+            when {
+                prop.nullable -> paramBuilder.defaultValue(CodeBlock.of("null"))
+                prop.defaultValue != null -> paramBuilder.defaultValue(formatDefaultValue(prop))
+            }
+
+            constructorBuilder.addParameter(paramBuilder.build())
+
+            PropertySpec
+                .builder(kotlinName, type)
+                .initializer(kotlinName)
+                .addAnnotation(AnnotationSpec.builder(SERIAL_NAME).addMember("%S", prop.name).build())
+                .build()
+        }
+
+        val typeSpecBuilder = TypeSpec
+            .classBuilder(variantSchema.name)
+            .addModifiers(KModifier.DATA)
+            .primaryConstructor(constructorBuilder.build())
+            .addProperties(propertySpecs)
+            .addAnnotation(SERIALIZABLE)
+            .addAnnotation(AnnotationSpec.builder(SERIAL_NAME).addMember("%S", serialName).build())
+            .superclass(parentClassName)
+            .addSuperclassConstructorParameter("")
+
+        if (variantSchema.description != null) {
+            typeSpecBuilder.addKdoc("%L", variantSchema.description)
+        }
+
+        return typeSpecBuilder.build()
+    }
+
     /**
      * Generates a JsonContentPolymorphicSerializer object for an anyOf schema without discriminator.
      */
@@ -285,12 +418,19 @@ class ModelGenerator(private val modelPackage: String) {
 
     /**
      * Generates a data class FileSpec, with superinterfaces and @SerialName resolved from hierarchy.
+     * Only anyOf-without-discriminator variants get superinterfaces here; oneOf/anyOf-with-discriminator
+     * variants are nested inside their parent sealed class by [generateSealedHierarchy].
      */
     context(hierarchy: HierarchyInfo)
     private fun generateDataClass(schema: SchemaModel): FileSpec {
         val className = ClassName(modelPackage, schema.name)
 
-        val parentEntries = hierarchy.variantParents[schema.name].orEmpty()
+        // Only apply superinterfaces for anyOf-without-discriminator variants
+        val parentEntries = hierarchy.variantParents[schema.name]
+            .orEmpty()
+            .filterKeys { parentClass ->
+                parentClass.simpleName in hierarchy.anyOfWithoutDiscriminator
+            }
         val serialName = parentEntries.values.firstOrNull()
         val superinterfaces = parentEntries.keys
 

core/src/main/kotlin/com/avsystem/justworks/core/gen/SerializersModuleGenerator.kt

@@ -37,7 +37,7 @@ class SerializersModuleGenerator(private val modelPackage: String) {
             val parentClass = ClassName(modelPackage, parent)
             code.beginControlFlow("%M(%T::class)", POLYMORPHIC_FUN, parentClass)
             for (variant in variants) {
-                val variantClass = ClassName(modelPackage, variant)
+                val variantClass = parentClass.nestedClass(variant)
                 code.addStatement("%M(%T::class)", SUBCLASS_FUN, variantClass)
             }
             code.endControlFlow()

core/src/main/kotlin/com/avsystem/justworks/core/gen/TypeMapping.kt

@@ -20,7 +20,11 @@ import com.squareup.kotlinpoet.TypeName
  * Maps [TypeRef] sealed variants to KotlinPoet [TypeName] instances.
  */
 object TypeMapping {
-    fun toTypeName(typeRef: TypeRef, modelPackage: String): TypeName = when (typeRef) {
+    fun toTypeName(
+        typeRef: TypeRef,
+        modelPackage: String,
+        classNameLookup: Map<String, ClassName> = emptyMap(),
+    ): TypeName = when (typeRef) {
         is TypeRef.Primitive -> {
             when (typeRef.type) {
                 PrimitiveType.STRING -> STRING
@@ -36,15 +40,15 @@ object TypeMapping {
         }
 
         is TypeRef.Array -> {
-            LIST.parameterizedBy(toTypeName(typeRef.items, modelPackage))
+            LIST.parameterizedBy(toTypeName(typeRef.items, modelPackage, classNameLookup))
         }
 
         is TypeRef.Map -> {
-            MAP.parameterizedBy(STRING, toTypeName(typeRef.valueType, modelPackage))
+            MAP.parameterizedBy(STRING, toTypeName(typeRef.valueType, modelPackage, classNameLookup))
         }
 
         is TypeRef.Reference -> {
-            ClassName(modelPackage, typeRef.schemaName)
+            classNameLookup[typeRef.schemaName] ?: ClassName(modelPackage, typeRef.schemaName)
         }
 
         is TypeRef.Inline -> {