/*
 * Copyright 2010-2023 JetBrains s.r.o. and Kotlin Programming Language contributors.
 * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
 */
package org.jetbrains.kotlinx.jspo.compiler.backend

import org.jetbrains.kotlin.backend.common.DeclarationTransformer
import org.jetbrains.kotlin.backend.common.compilationException
import org.jetbrains.kotlin.backend.common.extensions.IrGenerationExtension
import org.jetbrains.kotlin.backend.common.extensions.IrPluginContext
import org.jetbrains.kotlin.builtins.StandardNames
import org.jetbrains.kotlin.ir.UNDEFINED_OFFSET
import org.jetbrains.kotlin.ir.builders.declarations.buildFun
import org.jetbrains.kotlin.ir.declarations.*
import org.jetbrains.kotlin.ir.declarations.impl.IrVariableImpl
import org.jetbrains.kotlin.ir.expressions.IrBlockBody
import org.jetbrains.kotlin.ir.expressions.impl.*
import org.jetbrains.kotlin.ir.symbols.IrTypeParameterSymbol
import org.jetbrains.kotlin.ir.symbols.impl.IrVariableSymbolImpl
import org.jetbrains.kotlin.ir.types.IrType
import org.jetbrains.kotlin.ir.types.defaultType
import org.jetbrains.kotlin.ir.util.*
import org.jetbrains.kotlin.js.common.RESERVED_KEYWORDS
import org.jetbrains.kotlin.js.common.isES5IdentifierPart
import org.jetbrains.kotlin.js.common.isES5IdentifierStart
import org.jetbrains.kotlin.js.common.isValidES5Identifier
import org.jetbrains.kotlin.name.Name
import org.jetbrains.kotlin.util.OperatorNameConventions
import org.jetbrains.kotlinx.jspo.compiler.resolve.JsPlainObjectsPluginKey
import org.jetbrains.kotlinx.jspo.compiler.resolve.StandardIds
import kotlin.math.abs

/**
 * Because the inline members of external declarations are not participating in all the lowerings on the codegen phase,
 * we need to use this lowering to move their logic outside (in top-level inline functions) and just call those new proxy functions inside
 * the external inline members.
 *
 * So, the following generated by the plugin code:
 * ```
 * external interface Admin {
 *   val chat: Chat
 *   val email: String?
 *
 *   @JsExport.Ignore
 *   companion object {
 *      @JsNoDispatchReceiver
 *      inline operator fun invoke(chat: Chat, email: String? = VOID): Admin =
 *          js("{ chat: chat, name: name }")
 *
 *      @JsNoDispatchReceiver
 *      inline fun copy(source: Admin, chat: Chat = VOID, email: String = VOID): Admin =
 *          js("Object.assign({}, source, { chat: chat, email: email })")
 *   }
 * }
 * ```
 *
 * Will be translated in this one:
 * ```
 * inline fun Admin_invoke(chat: Chat, email: String?): Admin =
 *     js("{ chat: chat, name: name }")
 *
 * inline fun Admin_copy(source: Admin, chat: Chat, email: String?): Admin =
 *     js("Object.assign({}, source, { chat: chat, email: email })")
 *
 * external interface Admin {
 *   val chat: Chat
 *   val email: String?
 *
 *   @JsExport.Ignore
 *   companion object {
 *      @JsNoDispatchReceiver
 *      inline operator fun invoke(chat: Chat, email: String? = VOID): Admin =
 *          Admin_invoke(chat, email)
 *
 *      @JsNoDispatchReceiver
 *      inline fun copy(source: Admin, chat: Chat = VOID, email: String = VOID): Admin =
 *          Admin_copy(source, chat, email)
 *   }
 * }
 * ```
 */
private class MoveExternalInlineFunctionsWithBodiesOutsideLowering(private val context: IrPluginContext) : DeclarationTransformer {
    private val EXPECTED_ORIGIN = IrDeclarationOrigin.GeneratedByPlugin(JsPlainObjectsPluginKey)
    private val jsFunction = context.referenceFunctions(StandardIds.JS_FUNCTION_ID).single()

    override fun transformFlat(declaration: IrDeclaration): List<IrDeclaration>? {
        val file = declaration.file
        val parent = declaration.parentClassOrNull

        if (parent == null || declaration !is IrSimpleFunction || declaration.origin != EXPECTED_ORIGIN) return null

        val proxyFunction = createFunctionContainingTheLogic(declaration, parent).also(file::addChild)
        declaration.body = generateBodyWithTheProxyFunctionCall(declaration, proxyFunction)

        return null
    }

    private fun IrSimpleFunction.makeUniqueName(parent: IrClass): Name {
        val fqName = parent.fqNameWhenAvailable ?: compilationException("Parent class must have a FQ name", parent)
        val replacedWithUnderscoreFqn = sanitizeName(fqName.asString()).replace(".", "_")
        return Name.identifier(replacedWithUnderscoreFqn + "_" + sanitizeName(name.asString()))
    }

    private fun createFunctionContainingTheLogic(originalFunction: IrSimpleFunction, parent: IrClass): IrSimpleFunction {
        return context.irFactory.buildFun {
            updateFrom(originalFunction)
            name = originalFunction.makeUniqueName(parent)
            returnType = originalFunction.returnType
            origin = originalFunction.origin
            isInline = true
            isExternal = false
        }.apply {
            copyTypeParametersFrom(originalFunction)

            val substitutionMap = HashMap<IrTypeParameterSymbol, IrType>()
            substitutionMap.putAll(makeTypeParameterSubstitutionMap(originalFunction, this))

            copyParametersFrom(originalFunction, substitutionMap)

            returnType = returnType.substitute(substitutionMap)
            parameters = parameters.filter { it.kind != IrParameterKind.DispatchReceiver }

            body = when (originalFunction.name) {
                StandardNames.DATA_CLASS_COPY -> generateBodyForCopyFunction(this)
                OperatorNameConventions.INVOKE -> generateBodyForFactoryFunction(this)
                else -> compilationException("Unexpected function with name `${originalFunction.name.identifier}`", originalFunction)
            }
        }
    }

    private fun generateBodyWithTheProxyFunctionCall(originalFunction: IrSimpleFunction, proxyFunction: IrSimpleFunction): IrBlockBody {
        return context.irFactory.createBlockBody(originalFunction.startOffset, originalFunction.endOffset).apply {
            statements += IrReturnImpl(
                originalFunction.startOffset,
                originalFunction.endOffset,
                originalFunction.returnType,
                originalFunction.symbol,
                IrCallImpl(
                    originalFunction.startOffset,
                    originalFunction.endOffset,
                    originalFunction.returnType,
                    proxyFunction.symbol,
                    proxyFunction.typeParameters.size,
                ).apply {
                    for ((index, parameter) in originalFunction.parameters.withIndex()) {
                        if (parameter.kind != IrParameterKind.Regular) continue
                        arguments[index - 1] = IrGetValueImpl(UNDEFINED_OFFSET, UNDEFINED_OFFSET, parameter.symbol)
                    }

                    val typeParameters = originalFunction.typeParameters.ifEmpty { originalFunction.parentAsClass.typeParameters }
                    for ((index, type) in typeParameters.withIndex()) {
                        typeArguments[index] = type.defaultType
                    }
                }
            )
        }
    }

    private fun generateBodyForFactoryFunction(proxyFunction: IrSimpleFunction): IrBlockBody {
        return context.irFactory.createBlockBody(proxyFunction.startOffset, proxyFunction.endOffset).apply {
            statements += IrReturnImpl(
                proxyFunction.startOffset,
                proxyFunction.endOffset,
                proxyFunction.returnType,
                proxyFunction.symbol,
                IrCallImpl(
                    proxyFunction.startOffset,
                    proxyFunction.endOffset,
                    proxyFunction.returnType,
                    jsFunction,
                    0,
                ).apply {
                    val proxyFunctionRegularParameters = proxyFunction.parameters.filter { it.kind == IrParameterKind.Regular }
                    arguments[0] = createValueParametersObject(proxyFunctionRegularParameters).toIrConst(context.irBuiltIns.stringType)
                }
            )
        }
    }

    private fun sanitizeName(name: String): String {
        if (name.isEmpty()) return "_"

        val builder = StringBuilder(name.length + 7)

        val first = name.first()

        builder.append(first.mangleIfNot(Char::isES5IdentifierStart))

        for (idx in 1..name.lastIndex) {
            val c = name[idx]
            builder.append(c.mangleIfNot(Char::isES5IdentifierPart))
        }

        return "${builder}_${abs(name.hashCode()).toString(Character.MAX_RADIX)}"
    }

    private inline fun Char.mangleIfNot(predicate: Char.() -> Boolean) =
        if (predicate()) this else '_'

    private fun generateBodyForCopyFunction(proxyFunction: IrSimpleFunction): IrBlockBody {
        return context.irFactory.createBlockBody(proxyFunction.startOffset, proxyFunction.endOffset).apply {
            val selfName = proxyFunction.parameters.first().name
            statements += IrReturnImpl(
                proxyFunction.startOffset,
                proxyFunction.endOffset,
                proxyFunction.returnType,
                proxyFunction.symbol,
                IrCallImpl(
                    proxyFunction.startOffset,
                    proxyFunction.endOffset,
                    proxyFunction.returnType,
                    jsFunction,
                    0,
                ).apply {
                    val objectAssignCall =
                        "Object.assign({}, ${selfName.identifier}, ${createValueParametersObject(proxyFunction.parameters.drop(1))})"
                    arguments[0] = objectAssignCall.toIrConst(context.irBuiltIns.stringType)
                }
            )
        }
    }

    private fun createValueParametersObject(valueParameters: Iterable<IrValueParameter>): String {
        val listOfParameters = valueParameters.joinToString(", ") {
            val (key, value) = it.name.run {
                if (identifier.isValidES5Identifier() && identifier !in RESERVED_KEYWORDS) identifier to identifier
                else {
                    val newName: Name = Name.identifier(sanitizeName(identifier)).apply { it.name = this }
                    "\"${identifier}\"" to newName.identifier
                }
            }
            "$key:$value"
        }

        return "{$listOfParameters}"
    }
}

open class JsPlainObjectsLoweringExtension : IrGenerationExtension {
    override fun generate(moduleFragment: IrModuleFragment, pluginContext: IrPluginContext) {
        MoveExternalInlineFunctionsWithBodiesOutsideLowering(pluginContext).lower(moduleFragment)
    }
}
