threedom/threedom-kapt/src/main/kotlin/me/eater/threedom/kapt/EventNameProcessor.kt

package me.eater.threedom.kapt

import org.yanex.takenoko.*
import java.io.File
import javax.annotation.processing.*
import javax.lang.model.SourceVersion
import javax.lang.model.element.Element
import javax.lang.model.element.TypeElement
import javax.tools.Diagnostic.Kind.ERROR

@SupportedSourceVersion(SourceVersion.RELEASE_8)
@SupportedAnnotationTypes("me.eater.threedom.kapt.EventName")
@SupportedOptions(EventNameProcessor.KAPT_KOTLIN_GENERATED_OPTION_NAME, "org.gradle.annotation.processing.aggregating")
class EventNameProcessor : AbstractProcessor() {
    companion object {
        const val KAPT_KOTLIN_GENERATED_OPTION_NAME = "kapt.kotlin.generated"
    }

    override fun process(annotations: MutableSet<out TypeElement>?, roundEnv: RoundEnvironment): Boolean {
        val annotatedElements = roundEnv.getElementsAnnotatedWith(EventName::class.java)
        if (annotatedElements.isEmpty()) return false

        val kaptKotlinGeneratedDir = processingEnv.options[KAPT_KOTLIN_GENERATED_OPTION_NAME] ?: run {
            processingEnv.messager.printMessage(ERROR, "Can't find the target directory for generated Kotlin files.")
            return false
        }

        val generatedKtFile = kotlinFile("me.eater.threedom.generated") {
            objectDeclaration("EventNames") {
                val eventNames = mutableMapOf<String, String>()

                for (element in annotatedElements) {
                    val typeElement = element.toTypeElementOrNull() ?: continue
                    val eventName = typeElement.getAnnotation(EventName::class.java).eventName

                    if (eventName in eventNames) {
                        processingEnv.messager.printMessage(
                            ERROR,
                            "Class ${eventNames[eventName]} already uses the event name '${eventName}'",
                            element
                        )
                        continue
                    }

                    eventNames[typeElement.qualifiedName.toString()] = eventName

                    property(eventName) {
                        initializer(typeElement.qualifiedName.toString() + "::class")
                    }
                }

                property("EVENT_MAPPING") {
                    initializer(
                        "mapOf(\n${eventNames.map { (k, v) -> "    \"" + k.replace('"', '\"') + "\" to \"" + v.replace('"', '\"') + '"' }
                            .joinToString(",\n")}\n)"
                    )
                }

                function("getEventName") {
                    param<String>("eventClass")
                    returnType<String>()

                    body {
                        append("return EVENT_MAPPING[eventClass] ?: eventClass")
                    }
                }

                function(
                    "getEventName", INLINE
                ) {
                    typeParam("reified T")
                    returnType<String>()

                    body {
                        append("return getEventName(T::class.java.name)")
                    }
                }
            }
        }

        File("$kaptKotlinGeneratedDir/me/eater/threedom/generated", "EventNames.kt").apply {
            parentFile.mkdirs()
            writeText(generatedKtFile.accept(PrettyPrinter(PrettyPrinterConfiguration())))
        }

        return true
    }


    fun Element.toTypeElementOrNull(): TypeElement? {
        if (this !is TypeElement) {
            processingEnv.messager.printMessage(ERROR, "Invalid element type, class expected", this)
            return null
        }

        return this
    }
}