EnumGenerator.kt

internal class EnumGenerator(
  private val codeGenerator: CodeGenerator
) {

  fun generate(classDeclaration: KSClassDeclaration) {
    require(classDeclaration.classKind == ClassKind.ENUM_CLASS) {
      "${classDeclaration.containingFile?.filePath} is no enum."
    }

    val enum = buildEnum(classDeclaration)

    val fileSpec = FileSpec.builder(enum.name)
      .addType(enum)
      .build()

    createFile(classDeclaration, fileSpec)
  }

  private fun buildEnum(classDeclaration: KSClassDeclaration): TypeSpec {
    val fileName = classDeclaration.qualifiedName!!.asString()
      .removePrefix(classDeclaration.packageName.asString() + ".")
      .replace(".", "")

    val enumBuilder = TypeSpec.enumBuilder(fileName)
      .addModifiers(Modifier.PUBLIC)

    classDeclaration.declarations
      .forEach {
        if (it is KSClassDeclaration && it.classKind == ClassKind.ENUM_ENTRY) {
          val enumName = it.simpleName.asString().replaceFirstChar { char ->
            char.lowercase()
          }
          enumBuilder.addEnumCase(enumName)
        }
      }

    return enumBuilder.build()
  }

  private fun createFile(
    classDeclaration: KSClassDeclaration,
    fileSpec: FileSpec
  ) {
    codeGenerator.createNewFile(
      dependencies = Dependencies(
        aggregating = false,
        sources = listOfNotNull(classDeclaration.containingFile).toTypedArray()
      ),
      packageName = classDeclaration.packageName.asString(),
      extensionName = "swift",
      fileName = fileSpec.name
    ).use { output ->
      OutputStreamWriter(output, Charsets.UTF_8).use {
        fileSpec.writeTo(it)
      }
    }
  }
}