package viaduct.codegen.ct

import javassist.ClassPool
import javassist.CtClass
import javassist.bytecode.InnerClassesAttribute
import viaduct.codegen.utils.DEFAULT_IMPLS
import viaduct.codegen.utils.JavaBinaryName
import viaduct.codegen.utils.KmName

/**
 * This function rebuilds the InnerClassesAttribute ("ICA") of a class set to match what kotlinc produces.
 *
 * The high-level rule is that if class A makes *any reference* to nested class N, then N must have an entry
 * in the InnerClassesAttribute of class A.
 *
 * "Any reference" can include:
 *   - an actual nesting/nested relationship between A and N
 *   - N is a supertype of A
 *   - A uses N in a method or constructor body
 *   - N appears in a type parameter or annotation within A
 *
 * Since "any reference" is so all-encompassing, the rebuild process requires both a CtClass (to extract
 * references from function bodies) and a KmClassWrapper (to extract references from type parameters and
 * annotations that are not written to bytecode)
 */
fun rebuildInnerClassesAttribute(
    pool: ClassPool,
    results: Map<KmName, Pair<KmClassWrapper, CtClass>>
) = InnerClassesAttributeRebuilder(pool, results).rebuildAll()

private class InnerClassesAttributeRebuilder(
    private val pool: ClassPool,
    private val results: Map<KmName, Pair<KmClassWrapper, CtClass>>
) {
    private val seen = mutableSetOf<JavaBinaryName>()

    /** rebuild the [InnerClassesAttribute] of all classes in [results] */
    fun rebuildAll() {
        results.forEach { (kmName, v) ->
            val (_, ctClass) = v
            rebuild(kmName.asJavaBinaryName)

            // DefaultImpls are generated by the Ct layer rather than the Km layer, and won't
            // be included in the results map.
            // Make sure they get rebuilt too.
            ctClass.defaultImpls?.let { edge ->
                rebuild(edge.nested.javaBinaryName)
            }
        }
    }

    private fun rebuild(name: JavaBinaryName) {
        if (seen.contains(name)) return
        seen.add(name)

        if (!name.isGenerated) return
        val cls = name.ctClass ?: return

        val edges = mutableSetOf<NestEdge>()

        // for any references made by fields/behaviors of a class, if any of them are to
        // nested classes, then include them in this class's edges
        cls.refs.forEach {
            rebuild(it)
            buildEdges(edges, it.ctClass?.nestingEdge, false)
        }

        // javassist refs may not capture some kotlin-specific references to nested classes, such as
        // in annotations and type parameters.
        // Traverse through any kotlin-defined references to ensure that they are also captured.
        name.classWrapper?.let { wrapper ->
            wrapper.refs.map { it.asJavaBinaryName }.forEach { name ->
                rebuild(name)
                buildEdges(edges, name.ctClass?.nestingEdge, true)
            }
        }

        buildEdges(edges, cls.nestingEdge, includeSupers = true)

        if (edges.isNotEmpty()) {
            val ica = InnerClassesAttribute(cls.classFile.constPool)
            edges.forEach { it.write(ica) }

            // replace the existing InnerClassesAttribute
            cls.classFile.addAttribute(ica)
        }
    }

    private fun buildEdges(
        acc: MutableSet<NestEdge>,
        edge: NestEdge?,
        includeSupers: Boolean
    ) {
        if (edge != null && !acc.contains(edge)) {
            // write ancestral edge into nested class
            acc.add(edge)

            // if nested class A implement nested interfaces I, then A gets the lineage of
            // edges that nest I
            edge.nested.interfaces?.forEach {
                rebuild(it.javaBinaryName)
                if (includeSupers) {
                    // super edges are non-transitive; disable further
                    // super-traversal when adding this lineage
                    buildEdges(acc, it.nestingEdge, false)
                }
            }
            // if nested class A extends nested class C, then A gets the lineage of
            // edges that nest C
            edge.nested.superclass?.let {
                rebuild(it.javaBinaryName)
                if (includeSupers) {
                    // super edges are non-transitive; disable further
                    // super-traversal when adding this lineage
                    buildEdges(acc, it.nestingEdge, false)
                }
            }

            buildEdges(acc, edge.nesting.nestingEdge, includeSupers)
        }
    }

    /** if this class nests a DefaultImpls class, return that edge */
    private val CtClass.defaultImpls: NestEdge? get() =
        if (!this.isInterface) {
            null
        } else {
            this.nestedEdges.firstOrNull { it.nestedName == DEFAULT_IMPLS }
        }

    private val JavaBinaryName.ctClass: CtClass? get() = pool.getOrNull(this.toString())

    private val KmName.isGenerated: Boolean get() =
        results.containsKey(this) || this.toString().endsWith(".$DEFAULT_IMPLS")

    private val JavaBinaryName.isGenerated: Boolean get() = maybeAsKmName?.isGenerated ?: false

    private val JavaBinaryName.maybeAsKmName: KmName? get() =
        runCatching { asKmName }.getOrNull()

    private val JavaBinaryName.classWrapper: KmClassWrapper? get() =
        maybeAsKmName?.let { results[it]?.first }
}
