@file:Suppress("unused", "HasPlatformType")

package io.gitlab.jfronny.scripts

import com.squareup.javapoet.*
import org.gradle.api.Action
import org.gradle.api.Project
import org.gradle.api.tasks.SourceSet
import org.gradle.kotlin.dsl.extra
import java.lang.reflect.Type
import java.util.function.Supplier
import javax.lang.model.element.Element
import javax.lang.model.element.ExecutableElement
import javax.lang.model.element.Modifier
import javax.lang.model.type.DeclaredType
import javax.lang.model.type.TypeMirror
import javax.lang.model.util.Types

val Project.codegenDir get() = buildDir.resolve("generated/sources/jfCodegen")

fun SourceSet.generate(project: Project, generate: Action<ContentGenerator>) {
    val generators = project.extra["codeGenerators"] as LinkedHashMap<String, ContentGenerator.Generated>
    val exists = generators.containsKey(name)
    val generator = ContentGenerator().runAction(generate)
    generators[name] = if (exists) generators[name]!!.merge(generator.generate()) else generator.generate()
    if (!exists) {
        java.srcDir(project.codegenDir.resolve("java/$name"))
        resources.srcDir(project.codegenDir.resolve("resources/$name"))
    }
}

class ContentGenerator {
    private var finalized = false

    data class Generated(val classes: Map<String, String>, val resources: Map<String, ByteArray>) {
        fun merge(other: Generated): Generated = Generated(classes + other.classes, resources + other.resources)
    }

    private data class ImmutableMap<Key, Value>(private val inner: Map<Key, Value>): Map<Key, Value> by inner

    fun generate(): Generated {
        val fin = Generated(ImmutableMap(classes), ImmutableMap(resources))
        finalized = true
        return fin
    }

    private fun <T> ensureMutable(action: () -> T): T {
        if (finalized) throw IllegalAccessException("Attempted to access CodeGenerator after it was finalized")
        return action()
    }

    private val classes: MutableMap<String, String> = LinkedHashMap()
    private val resources: MutableMap<String, ByteArray> = LinkedHashMap()

    fun `class`(`package`: String, name: String, generator: Action<TypeSpec.Builder>) = `class`(`package`, name, TypeSpec::classBuilder, generator)
    fun `enum`(`package`: String, name: String, generator: Action<TypeSpec.Builder>) = `class`(`package`, name, TypeSpec::enumBuilder, generator)
    fun `interface`(`package`: String, name: String, generator: Action<TypeSpec.Builder>) = `class`(`package`, name, TypeSpec::interfaceBuilder, generator)

    private fun `class`(`package`: String, name: String, builder: (ClassName) -> TypeSpec.Builder, generator: Action<TypeSpec.Builder>) = ensureMutable {
        check(`package`, "package", packagePattern)
        check(name, "class name", classNamePattern)
        val bld = builder(ClassName.get(`package`, name))
        generator.execute(bld)
        val javaFile: JavaFile = JavaFile.builder(`package`, bld.build())
            .skipJavaLangImports(true)
            .indent("    ")
            .addFileComment("Automatically generated through JfCodegen, do not edit!")
            .build()
        classes["${`package`.replace('.', '/')}/$name.java"] = javaFile.toString()
    }

    fun text(path: String, generator: Action<StringBuilder>) = blob(path) {
        val builder = StringBuilder()
        generator.execute(builder)
        builder.toString().toByteArray()
    }

    fun blob(path: String, generator: Supplier<ByteArray>) = ensureMutable {
        check(path, "path", pathPattern)
        resources[path] = generator.get()
    }

    companion object {
        @JvmStatic private val packagePattern = Regex("(?:[a-z][a-z0-9_]*(\\.[a-z][a-z0-9_]*)*)")
        @JvmStatic private val classNamePattern = Regex("(?:[A-Z][A-Za-z0-9_\$]*)")
        @JvmStatic private val pathPattern = Regex("(?:([a-z]+/)*[a-zA-Z][a-zA-Z0-9_.]*)")

        @JvmStatic
        private fun check(name: String, type: String, pattern: Regex) {
            if (!pattern.matches(name)) throw IllegalArgumentException("$type \"$name\" is not a valid $type")
        }
    }
}

// Extensions for TypeSpec.Builder
fun TypeSpec.Builder.javadoc(`import`: String) = addJavadoc(import)
fun TypeSpec.Builder.javadoc(generate: Action<CodeBlock.Builder>) = addJavadoc(CodeBlock.builder().runAction(generate).build())
fun TypeSpec.Builder.annotation(annotation: Class<*>, generate: Action<AnnotationSpec.Builder>? = null) = addAnnotation(AnnotationSpec.builder(annotation).runAction(generate).build())
fun TypeSpec.Builder.annotation(annotation: ClassName, generate: Action<AnnotationSpec.Builder>? = null) = addAnnotation(AnnotationSpec.builder(annotation).runAction(generate).build())
fun TypeSpec.Builder.modifiers(vararg modifiers: Modifier) = addModifiers(*modifiers)
fun TypeSpec.Builder.typeVariable(typeVariable: TypeVariableName) = addTypeVariable(typeVariable)
fun TypeSpec.Builder.superInterface(superInterface: TypeName) = addSuperinterface(superInterface)
fun TypeSpec.Builder.superInterface(superInterface: Type, avoidNestedTypeNameClashes: Boolean = true) = addSuperinterface(superInterface, avoidNestedTypeNameClashes)
fun TypeSpec.Builder.superInterface(superInterface: TypeMirror, avoidNestedTypeNameClashes: Boolean = true) = addSuperinterface(superInterface, avoidNestedTypeNameClashes)
fun TypeSpec.Builder.enumConstant(name: String, typeSpec: TypeSpec? = null) = if (typeSpec == null) addEnumConstant(name) else addEnumConstant(name, typeSpec)
fun TypeSpec.Builder.field(type: Type, name: String, vararg modifiers: Modifier, generate: Action<FieldSpec.Builder>? = null) = addField(FieldSpec.builder(type, name, *modifiers).runAction(generate).build())
fun TypeSpec.Builder.field(type: TypeName, name: String, vararg modifiers: Modifier, generate: Action<FieldSpec.Builder>? = null) = addField(FieldSpec.builder(type, name, *modifiers).runAction(generate).build())
fun TypeSpec.Builder.field(name: String, value: Boolean, vararg modifiers: Modifier) = field(Boolean::class.java, name, *modifiers) { initializer("$value") }
fun TypeSpec.Builder.field(name: String, value: Short, vararg modifiers: Modifier) = field(Short::class.java, name, *modifiers) { initializer("$value") }
fun TypeSpec.Builder.field(name: String, value: Int, vararg modifiers: Modifier) = field(Int::class.java, name, *modifiers) { initializer("$value") }
fun TypeSpec.Builder.field(name: String, value: Long, vararg modifiers: Modifier) = field(Long::class.java, name, *modifiers) { initializer("${value}L") }
fun TypeSpec.Builder.field(name: String, value: Float, vararg modifiers: Modifier) = field(Float::class.java, name, *modifiers) { initializer("${value}f") }
fun TypeSpec.Builder.field(name: String, value: Double, vararg modifiers: Modifier) = field(Double::class.java, name, *modifiers) { initializer("$value") }
fun TypeSpec.Builder.field(name: String, value: String, vararg modifiers: Modifier) = field(String::class.java, name, *modifiers) { initializer("\$S", value) }
fun TypeSpec.Builder.static(generate: Action<CodeBlock.Builder>) = addStaticBlock(CodeBlock.builder().runAction(generate).build())
fun TypeSpec.Builder.initializer(generate: Action<CodeBlock.Builder>) = addInitializerBlock(CodeBlock.builder().runAction(generate).build())
fun TypeSpec.Builder.method(name: String, generate: Action<MethodSpec.Builder>) = addMethod(MethodSpec.methodBuilder(name).runAction(generate).build())
fun TypeSpec.Builder.override(element: ExecutableElement, generate: Action<MethodSpec.Builder>) = addMethod(MethodSpec.overriding(element).runAction(generate).build())
fun TypeSpec.Builder.override(element: ExecutableElement, enclosing: DeclaredType, types: Types, generate: Action<MethodSpec.Builder>) = addMethod(MethodSpec.overriding(element, enclosing, types).runAction(generate).build())
fun TypeSpec.Builder.constructor(generate: Action<MethodSpec.Builder>) = addMethod(MethodSpec.constructorBuilder().runAction(generate).build())
fun TypeSpec.Builder.innerEnum(name: String, generate: Action<TypeSpec.Builder>) = addType(TypeSpec.enumBuilder(name).runAction(generate).build())
fun TypeSpec.Builder.innerInterface(name: String, generate: Action<TypeSpec.Builder>) = addType(TypeSpec.interfaceBuilder(name).runAction(generate).build())
fun TypeSpec.Builder.innerClass(name: String, generate: Action<TypeSpec.Builder>) = addType(TypeSpec.classBuilder(name).runAction(generate).build())
fun TypeSpec.Builder.innerAnnotation(name: String, generate: Action<TypeSpec.Builder>) = addType(TypeSpec.annotationBuilder(name).runAction(generate).build())
fun TypeSpec.Builder.innerEnum(name: ClassName, generate: Action<TypeSpec.Builder>) = addType(TypeSpec.enumBuilder(name).runAction(generate).build())
fun TypeSpec.Builder.innerInterface(name: ClassName, generate: Action<TypeSpec.Builder>) = addType(TypeSpec.interfaceBuilder(name).runAction(generate).build())
fun TypeSpec.Builder.innerClass(name: ClassName, generate: Action<TypeSpec.Builder>) = addType(TypeSpec.classBuilder(name).runAction(generate).build())
fun TypeSpec.Builder.innerAnnotation(name: ClassName, generate: Action<TypeSpec.Builder>) = addType(TypeSpec.annotationBuilder(name).runAction(generate).build())
fun TypeSpec.Builder.origin(origin: Element) = addOriginatingElement(origin)

// Extensions for AnnotationSpec.Builder
fun AnnotationSpec.Builder.member(name: String, format: String, vararg args: Any) = addMember(name, format, args)
fun AnnotationSpec.Builder.member(name: String, generate: Action<CodeBlock.Builder>) = addMember(name, CodeBlock.builder().runAction(generate).build())

// Extensions for FieldSpec.Builder
fun FieldSpec.Builder.initializer(generate: Action<CodeBlock.Builder>) = initializer(CodeBlock.builder().runAction(generate).build())

// Extensions for MethodSpec.Builder
fun MethodSpec.Builder.javadoc(`import`: String) = addJavadoc(import)
fun MethodSpec.Builder.javadoc(generate: Action<CodeBlock.Builder>) = addJavadoc(CodeBlock.builder().runAction(generate).build())
fun MethodSpec.Builder.annotation(annotation: Class<*>, generate: Action<AnnotationSpec.Builder>? = null) = addAnnotation(AnnotationSpec.builder(annotation).runAction(generate).build())
fun MethodSpec.Builder.annotation(annotation: ClassName, generate: Action<AnnotationSpec.Builder>? = null) = addAnnotation(AnnotationSpec.builder(annotation).runAction(generate).build())
fun MethodSpec.Builder.modifiers(vararg modifiers: Modifier) = addModifiers(*modifiers)
fun MethodSpec.Builder.typeVariable(typeVariable: TypeVariableName) = addTypeVariable(typeVariable)
fun MethodSpec.Builder.parameter(type: Type, name: String, vararg modifiers: Modifier, generate: Action<ParameterSpec.Builder>? = null) = addParameter(ParameterSpec.builder(type, name, *modifiers).runAction(generate).build())
fun MethodSpec.Builder.parameter(type: TypeName, name: String, vararg modifiers: Modifier, generate: Action<ParameterSpec.Builder>? = null) = addParameter(ParameterSpec.builder(type, name, *modifiers).runAction(generate).build())
fun MethodSpec.Builder.exception(exception: Type) = addException(exception)
fun MethodSpec.Builder.exception(exception: TypeName) = addException(exception)
fun MethodSpec.Builder.code(generate: Action<CodeBlock.Builder>) = addCode(CodeBlock.builder().runAction(generate).build())
fun MethodSpec.Builder.defaultValue(generate: Action<CodeBlock.Builder>) = defaultValue(CodeBlock.builder().runAction(generate).build())

// Extensions for ParameterSpec.Builder
fun ParameterSpec.Builder.javadoc(`import`: String) = addJavadoc(import)
fun ParameterSpec.Builder.javadoc(generate: Action<CodeBlock.Builder>) = addJavadoc(CodeBlock.builder().runAction(generate).build())
fun ParameterSpec.Builder.annotation(annotation: Class<*>, generate: Action<AnnotationSpec.Builder>? = null) = addAnnotation(AnnotationSpec.builder(annotation).runAction(generate).build())
fun ParameterSpec.Builder.annotation(annotation: ClassName, generate: Action<AnnotationSpec.Builder>? = null) = addAnnotation(AnnotationSpec.builder(annotation).runAction(generate).build())
fun ParameterSpec.Builder.modifiers(vararg modifiers: Modifier) = addModifiers(*modifiers)

// Extensions for CodeBlock.Builder
fun CodeBlock.Builder.controlFlow(controlFlow: String, vararg args: Any, generate: Action<CodeBlock.Builder>) = beginControlFlow(controlFlow, *args).runAction(generate).endControlFlow()
fun CodeBlock.Builder.indent(generate: Action<CodeBlock.Builder>) = indent().runAction(generate).unindent()