#!/usr/bin/env kotlin

@file:DependsOn("org.jetbrains.kotlin:kotlin-compiler-embeddable:2.3.21")

import org.jetbrains.kotlin.cli.jvm.compiler.KotlinCoreApplicationEnvironment
import org.jetbrains.kotlin.cli.jvm.compiler.KotlinCoreApplicationEnvironmentMode
import org.jetbrains.kotlin.cli.jvm.compiler.KotlinCoreProjectEnvironment
import org.jetbrains.kotlin.com.intellij.openapi.util.Disposer
import org.jetbrains.kotlin.parsing.KotlinParserDefinition
import org.jetbrains.kotlin.psi.KtClass
import org.jetbrains.kotlin.psi.KtPsiFactory
import java.nio.file.Files
import java.nio.file.Path
import kotlin.io.path.*
import kotlin.system.exitProcess


fun process(file: Path, outputDir: Path, factory: KtPsiFactory) {
    if(!file.exists()) {
        System.err.println("Skipping $file: file does not exist")
        return
    }
    if(!file.isReadable()) {
        System.err.println("Skipping $file: file is not readable")
        return
    }

    val ktFile = factory.createFile(file.readText())
    val pkg = ktFile.packageFqName.toString()
    val imports = ktFile.importDirectives.mapNotNull { it.importedFqName }.filter { !it.toString().contains("ToPartialize") }
    val classes = ktFile.declarations.mapNotNull { cls ->
        if(cls !is KtClass || !cls.isData()) return@mapNotNull null
        val annotations = cls.annotationEntries.map { it.text }
        if("@ToPartialize" !in annotations) return@mapNotNull null
        val ogName = cls.name ?: return@mapNotNull null

        val newName = "Partial$ogName"

        val nestedTypes = cls.declarations.filterIsInstance<KtClass>().mapNotNull { it.name }.toSet()

        val props = cls.primaryConstructorParameters.map {
            val name = it.name
            val type = it.typeReference?.text
            if(name == null || type == null) return@mapNotNull null
            name to type
        }.joinToString(", ") { (n, t) ->
            val nnT = if(t.endsWith('?')) t.substring(startIndex = 0, endIndex = t.length - 1) else t
            val useT = if(nnT in nestedTypes) "$ogName.$nnT" else nnT
            "val $n: $useT? = null"
        }

        "@Serializable\ndata class $newName($props)"
    }

    val cnt = "package $pkg\n\n${imports.joinToString("\n") { "import $it" }}\n\n${classes.joinToString("\n\n")}"
//    println(cnt)

    val writeDir = outputDir.resolve(pkg.replace('.', '/'))
    val fileName = file.fileName
//    println("Trying to write to $writeDir/$fileName")

    if(!writeDir.exists()) writeDir.createDirectories()
    Files.writeString(Path("$writeDir/$fileName"), cnt)
}


// usage kotlin partialize.main.kts output-dir [input-file]+
if(args.isEmpty()) {
    System.err.println("Usage: partialize.main.kts <output-dir> <input-file>+")
    exitProcess(-1)
}

val outputDir = Path(args[0])
if(outputDir.notExists()) {
    println("Creating output directory $outputDir...")
    Files.createDirectory(outputDir)
}
else if(!outputDir.isDirectory()) {
    System.err.println("Output directory $outputDir is not a directory.")
    exitProcess(-1)
}
else if(!outputDir.isWritable()) {
    System.err.println("Cannot write to output directory $outputDir")
    exitProcess(-1)
}

val inputFiles = args.slice(1 until args.size)
val disp = Disposer.newDisposable()
val appEnv = KotlinCoreApplicationEnvironment.create(disp, KotlinCoreApplicationEnvironmentMode.Production)
appEnv.registerParserDefinition(KotlinParserDefinition())
val project = KotlinCoreProjectEnvironment(disp, appEnv)
val factory = KtPsiFactory(project.project)
inputFiles.forEach {
    process(Path(it), outputDir, factory)
}