/*
 * Copyright 2010-2021 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.kotlin.backend.wasm.lower

import org.jetbrains.kotlin.backend.common.BodyLoweringPass
import org.jetbrains.kotlin.backend.common.DeclarationTransformer
import org.jetbrains.kotlin.backend.common.ir.addDispatchReceiver
import org.jetbrains.kotlin.backend.common.lower.DeclarationIrBuilder
import org.jetbrains.kotlin.backend.common.lower.createIrBuilder
import org.jetbrains.kotlin.backend.wasm.WasmBackendContext
import org.jetbrains.kotlin.backend.wasm.ir2wasm.isBuiltInWasmRefType
import org.jetbrains.kotlin.backend.wasm.ir2wasm.isExported
import org.jetbrains.kotlin.backend.wasm.ir2wasm.isExternalType
import org.jetbrains.kotlin.descriptors.DescriptorVisibilities
import org.jetbrains.kotlin.ir.backend.js.utils.getJsNameOrKotlinName
import org.jetbrains.kotlin.ir.builders.*
import org.jetbrains.kotlin.ir.builders.declarations.*
import org.jetbrains.kotlin.ir.declarations.*
import org.jetbrains.kotlin.ir.expressions.IrBody
import org.jetbrains.kotlin.ir.expressions.IrCall
import org.jetbrains.kotlin.ir.expressions.IrExpression
import org.jetbrains.kotlin.ir.expressions.IrStatementOriginImpl
import org.jetbrains.kotlin.ir.expressions.impl.IrInstanceInitializerCallImpl
import org.jetbrains.kotlin.ir.types.*
import org.jetbrains.kotlin.ir.util.*
import org.jetbrains.kotlin.ir.visitors.IrElementTransformerVoid
import org.jetbrains.kotlin.ir.visitors.transformChildrenVoid
import org.jetbrains.kotlin.name.Name
import kotlin.math.absoluteValue

/**
 * Create wrappers for external and @JsExport functions when type adaptation is needed
 */
class JsInteropFunctionsLowering(val context: WasmBackendContext) : DeclarationTransformer {
    val builtIns = context.irBuiltIns
    val symbols = context.wasmSymbols
    val adapters = symbols.jsInteropAdapters

    // Used to for export lambdas
    object KOTLIN_WASM_CLOSURE_FOR_JS_CLOSURE : IrStatementOriginImpl("KOTLIN_WASM_CLOSURE_FOR_JS_CLOSURE")

    private val closureCallExports = mutableMapOf<IrSimpleType, IrSimpleFunction>()
    private val kotlinClosureToJsConverters = mutableMapOf<IrSimpleType, IrSimpleFunction>()
    private val jsClosureCallers = mutableMapOf<IrSimpleType, IrSimpleFunction>()
    private val jsToKotlinClosures = mutableMapOf<IrSimpleType, IrSimpleFunction>()

    val additionalDeclarations = mutableListOf<IrDeclaration>()
    lateinit var currentParent: IrDeclarationParent

    override fun transformFlat(declaration: IrDeclaration): List<IrDeclaration>? {
        if (declaration.isFakeOverride) return null
        if (declaration !is IrSimpleFunction) return null
        val isExported = declaration.isExported()
        val isExternal = declaration.isExternal
        if (declaration.isPropertyAccessor) return null
        if (declaration.parent !is IrPackageFragment) return null
        if (!isExported && !isExternal) return null
        check(!(isExported && isExternal)) { "Exported external declarations are not supported: ${declaration.fqNameWhenAvailable}" }
        check(declaration.parent !is IrClass) { "Interop members are not supported:  ${declaration.fqNameWhenAvailable}" }

        additionalDeclarations.clear()
        currentParent = declaration.parent
        val newDeclarations = if (isExternal)
            transformExternalFunction(declaration)
        else
            transformExportFunction(declaration)

        return (newDeclarations ?: listOf(declaration)) + additionalDeclarations
    }

    /**
     *  external fun foo(x: KotlinType): KotlinType
     *
     *  ->
     *
     *  external fun foo(x: JsType): JsType
     *  fun foo__externalAdapter(x: KotlinType): KotlinType = adaptResult(foo(adaptParameter(x)));
     */
    fun transformExternalFunction(function: IrSimpleFunction): List<IrDeclaration>? {
        // External functions with default parameter values are already processed by
        // [ComplexExternalDeclarationsToTopLevelFunctionsLowering]
        if (function.valueParameters.any { it.defaultValue != null })
            return null

        val valueParametersAdapters = function.valueParameters.map {
            it.type.kotlinToJsAdapterIfNeeded(isReturn = false)
        }
        val resultAdapter =
            function.returnType.jsToKotlinAdapterIfNeeded(isReturn = true)

        if (resultAdapter == null && valueParametersAdapters.all { it == null })
            return null

        val newFun = context.irFactory.createStaticFunctionWithReceivers(
            function.parent,
            name = Name.identifier(function.name.asStringStripSpecialMarkers() + "__externalAdapter"),
            function
        )

        function.valueParameters.forEachIndexed { index, newParameter ->
            val adapter = valueParametersAdapters[index]
            if (adapter != null) {
                newParameter.type = adapter.toType
            }
        }
        resultAdapter?.let {
            function.returnType = resultAdapter.fromType
        }

        val builder = context.createIrBuilder(newFun.symbol)
        newFun.body = createAdapterFunctionBody(builder, newFun, function, valueParametersAdapters, resultAdapter)
        newFun.annotations = emptyList()

        context.mapping.wasmJsInteropFunctionToWrapper[function] = newFun
        return listOf(function, newFun)
    }

    /**
     *  @JsExport
     *  fun foo(x: KotlinType): KotlinType { <original-body> }
     *
     *  ->
     *
     *  @JsExport
     *  @JsName("foo")
     *  fun foo__JsExportAdapter(x: JsType): JsType =
     *      adaptResult(foo(adaptParameter(x)));
     *
     *  fun foo(x: KotlinType): KotlinType { <original-body> }
     */
    fun transformExportFunction(function: IrSimpleFunction): List<IrDeclaration>? {
        val valueParametersAdapters = function.valueParameters.map {
            it.type.jsToKotlinAdapterIfNeeded(isReturn = false)
        }
        val resultAdapter =
            function.returnType.kotlinToJsAdapterIfNeeded(isReturn = true)

        if (resultAdapter == null && valueParametersAdapters.all { it == null })
            return null

        val newFun = context.irFactory.createStaticFunctionWithReceivers(
            function.parent,
            name = Name.identifier(function.name.asStringStripSpecialMarkers() + "__JsExportAdapter"),
            function
        )

        newFun.valueParameters.forEachIndexed { index, newParameter ->
            val adapter = valueParametersAdapters[index]
            if (adapter != null) {
                newParameter.type = adapter.fromType
            }
        }
        resultAdapter?.let {
            newFun.returnType = resultAdapter.toType
        }

        // Delegate new function to old function:
        val builder: DeclarationIrBuilder = context.createIrBuilder(newFun.symbol)
        newFun.body = createAdapterFunctionBody(builder, newFun, function, valueParametersAdapters, resultAdapter)

        newFun.annotations += builder.irCallConstructor(context.wasmSymbols.jsNameConstructor, typeArguments = emptyList()).also {
            it.putValueArgument(0, builder.irString(function.getJsNameOrKotlinName().identifier))
        }
        function.annotations = function.annotations.filter { it.symbol != context.wasmSymbols.jsExportConstructor }

        return listOf(function, newFun)
    }

    private fun createAdapterFunctionBody(
        builder: DeclarationIrBuilder,
        function: IrSimpleFunction,
        functionToCall: IrSimpleFunction,
        valueParametersAdapters: List<InteropTypeAdapter?>,
        resultAdapter: InteropTypeAdapter?
    ) = builder.irBlockBody {
        +irReturn(
            irCall(functionToCall).let { call ->
                for ((index, valueParameter) in function.valueParameters.withIndex()) {
                    val get = irGet(valueParameter)
                    call.putValueArgument(index, valueParametersAdapters[index].adaptIfNeeded(get, builder))
                }
                resultAdapter.adaptIfNeeded(call, builder)
            }
        )
    }

    private fun IrType.kotlinToJsAdapterIfNeeded(isReturn: Boolean): InteropTypeAdapter? {
        if (isReturn && this == builtIns.unitType)
            return null

        when (this) {
            builtIns.stringType -> return FunctionBasedAdapter(adapters.kotlinToJsStringAdapter.owner)
            builtIns.stringType.makeNullable() -> return NullOrAdapter(FunctionBasedAdapter(adapters.kotlinToJsStringAdapter.owner))
            builtIns.booleanType -> return FunctionBasedAdapter(adapters.kotlinToJsBooleanAdapter.owner)
            builtIns.anyType -> return FunctionBasedAdapter(adapters.kotlinToJsAnyAdapter.owner)

            builtIns.byteType,
            builtIns.shortType,
            builtIns.charType,
            builtIns.intType,
            builtIns.longType,
            builtIns.floatType,
            builtIns.doubleType,
            context.wasmSymbols.voidType ->
                return null

        }

        if (isExternalType(this))
            return null

        if (isBuiltInWasmRefType(this))
            return null

        if (this is IrSimpleType && this.isFunction()) {
            val functionTypeInfo = FunctionTypeInfo(this, toJs = true)

            // Kotlin's closures are objects that implement FunctionN interface.
            // JavaScript can receive opaque reference to them but cannot call them directly.
            // Thus, we export helper "caller" method that JavaScript will use to call kotlin closures:
            //
            //     @JsExport
            //     fun __callFunction_<signatureHash>(f: dataref, p1: JsType1, p2: JsType2, ...): JsTypeRes {
            //          return adapt(
            //              cast<FunctionN>(f).invoke(
            //                  adapt(p1),
            //                  adapt(p2),
            //                  ...
            //               )
            //          )
            //     }
            //
            closureCallExports.getOrPut(this) {
                createKotlinClosureCaller(functionTypeInfo)
            }

            // Converter functions creates new JavaScript closures that delegate to Kotlin closures
            // using above-mentioned "caller" export:
            //
            //     @JsFun("""(f) => {
            //        (p1, p2, ...) => <wasm-exports>.__callFunction_<signatureHash>(f, p1, p2, ...)
            //     }""")
            //     external fun __convertKotlinClosureToJsClosure_<signatureHash>(f: dataref): ExternalRef
            //
            val kotlinToJsClosureConvertor = kotlinClosureToJsConverters.getOrPut(this) {
                createKotlinToJsClosureConvertor(functionTypeInfo)
            }
            return FunctionBasedAdapter(kotlinToJsClosureConvertor)
        }

        return SendKotlinObjectToJsAdapter(this)
    }

    private fun IrType.jsToKotlinAdapterIfNeeded(isReturn: Boolean): InteropTypeAdapter? {
        if (isReturn && this == builtIns.unitType)
            return null

        when (this) {
            builtIns.stringType -> return FunctionBasedAdapter(adapters.jsToKotlinStringAdapter.owner)
            builtIns.anyType -> return FunctionBasedAdapter(adapters.jsToKotlinAnyAdapter.owner)
            builtIns.byteType -> return FunctionBasedAdapter(adapters.jsToKotlinByteAdapter.owner)
            builtIns.shortType -> return FunctionBasedAdapter(adapters.jsToKotlinShortAdapter.owner)
            builtIns.charType -> return FunctionBasedAdapter(adapters.jsToKotlinCharAdapter.owner)

            builtIns.booleanType,
            builtIns.intType,
            builtIns.longType,
            builtIns.floatType,
            builtIns.doubleType,
            context.wasmSymbols.voidType ->
                return null
        }

        if (isExternalType(this))
            return null

        if (isBuiltInWasmRefType(this))
            return null

        if (this is IrSimpleType && this.isFunction()) {
            val functionTypeInfo = FunctionTypeInfo(this, toJs = false)

            // JavaScript's closures are external references that cannot be called directly in WebAssembly.
            // Thus, we import helper "caller" method that WebAssembly will use to call JS closures:
            //
            //     @JsFun("(f, p0, p1, ...) => f(p0, p1, ...)")
            //     external fun __callJsClosure_<signatureHash>(f: ExternalRef, p0: JsType1, p1: JsType2, ...): JsResType
            //
            val jsClosureCaller = jsClosureCallers.getOrPut(this) {
                createJsClosureCaller(functionTypeInfo)
            }

            // Converter functions creates new Kotlin closure that delegate to JS closure
            // using above-mentioned "caller" import:
            //
            //     fun __convertJsClosureToKotlinClosure_<signatureHash>(f: ExternalRef) : FunctionN<KotlinType1, ..., KotlinResType> =
            //       { p0: KotlinType1, p1: KotlinType2, ... ->
            //          adapt(__callJsClosure_<signatureHash>(f, adapt(p0), adapt(p1), ..))
            //       }
            //
            val jsToKotlinClosure = jsToKotlinClosures.getOrPut(this) {
                createJsToKotlinClosureConverter(functionTypeInfo, jsClosureCaller)
            }
            return FunctionBasedAdapter(jsToKotlinClosure)
        }

        return ReceivingKotlinObjectFromJsAdapter(this)
    }

    private fun createKotlinClosureCaller(info: FunctionTypeInfo): IrSimpleFunction {
        val result = context.irFactory.buildFun {
            name = Name.identifier("__callFunction_${info.hashString}")
            returnType = info.adaptedResultType
        }
        result.parent = currentParent
        result.addValueParameter {
            name = Name.identifier("f")
            type = context.wasmSymbols.wasmDataRefType
        }
        var count = 0
        info.adaptedParameterTypes.forEach { type ->
            result.addValueParameter {
                this.name = Name.identifier("p" + count++.toString())
                this.type = type
            }
        }
        val builder = context.createIrBuilder(result.symbol)

        result.body = builder.irBlockBody {
            val invokeFun = info.functionType.classOrNull!!.owner.functions.single { it.name == Name.identifier("invoke") }
            val callInvoke = irCall(invokeFun.symbol, info.originalResultType).also { call ->
                call.dispatchReceiver =
                    ReceivingKotlinObjectFromJsAdapter(invokeFun.dispatchReceiverParameter!!.type)
                        .adapt(irGet(result.valueParameters[0]), builder)

                for (i in info.adaptedParameterTypes.indices) {
                    call.putValueArgument(i, info.parametersAdapters[i].adaptIfNeeded(irGet(result.valueParameters[i + 1]), builder))
                }
            }
            +irReturn(info.resultAdapter.adaptIfNeeded(callInvoke, builder))
        }
        result.annotations += builder.irCallConstructor(context.wasmSymbols.jsExportConstructor, typeArguments = emptyList())
        additionalDeclarations += result
        return result
    }

    private fun createKotlinToJsClosureConvertor(info: FunctionTypeInfo): IrSimpleFunction {
        val result = context.irFactory.buildFun {
            name = Name.identifier("__convertKotlinClosureToJsClosure_${info.hashString}")
            returnType = context.wasmSymbols.externalInterfaceType
        }
        result.parent = currentParent
        result.addValueParameter {
            name = Name.identifier("f")
            type = context.wasmSymbols.wasmDataRefType
        }
        val builder = context.createIrBuilder(result.symbol)
        // TODO: Cache created JS closures
        val arity = info.parametersAdapters.size
        val jsCode = buildString {
            append("(f) => (")
            appendParameterList(arity)
            append(") => wasmExports.__callFunction_")
            append(info.hashString)
            append("(f, ")
            appendParameterList(arity)
            append(")")
        }

        result.annotations += builder.irCallConstructor(context.wasmSymbols.jsFunConstructor, typeArguments = emptyList()).also {
            it.putValueArgument(0, builder.irString(jsCode))
        }

        additionalDeclarations += result
        return result
    }

    private fun createJsToKotlinClosureConverter(
        info: FunctionTypeInfo,
        jsClosureCaller: IrSimpleFunction,
    ): IrSimpleFunction {
        val functionType = info.functionType
        val result = context.irFactory.buildFun {
            name = Name.identifier("__convertJsClosureToKotlinClosure_${info.hashString}")
            returnType = functionType
        }
        result.parent = currentParent
        result.addValueParameter {
            name = Name.identifier("f")
            type = context.wasmSymbols.externalInterfaceType
        }

        val closureClass = context.irFactory.buildClass {
            name = Name.identifier("__JsClosureToKotlinClosure_${info.hashString}")
        }.apply {
            createImplicitParameterDeclarationWithWrappedDescriptor()
            superTypes = listOf(functionType)
            parent = currentParent
        }

        val closureClassField = closureClass.addField {
            name = Name.identifier("jsClosure")
            type = context.wasmSymbols.externalInterfaceType
            visibility = DescriptorVisibilities.PRIVATE
            isFinal = true
        }

        val closureClassConstructor = closureClass.addConstructor {
            isPrimary = true
        }.apply {
            val parameter = addValueParameter {
                name = closureClassField.name
                type = closureClassField.type
            }
            body = context.createIrBuilder(symbol).irBlockBody(startOffset, endOffset) {
                +irDelegatingConstructorCall(context.irBuiltIns.anyClass.owner.constructors.single())
                +irSetField(irGet(closureClass.thisReceiver!!), closureClassField, irGet(parameter))
                +IrInstanceInitializerCallImpl(startOffset, endOffset, closureClass.symbol, context.irBuiltIns.unitType)
            }
        }

        closureClass.addFunction {
            name = Name.identifier("invoke")
            returnType = info.originalResultType
        }.apply {
            addDispatchReceiver { type = closureClass.defaultType }
            info.originalParameterTypes.forEachIndexed { index, irType ->
                addValueParameter {
                    name = Name.identifier("p$index")
                    type = irType
                }
            }
            val lambdaBuilder = context.createIrBuilder(symbol)
            body = lambdaBuilder.irBlockBody {
                val jsClosureCallerCall = irCall(jsClosureCaller)
                jsClosureCallerCall.putValueArgument(0, irGetField(irGet(dispatchReceiverParameter!!), closureClassField))
                for ((adapterIndex, paramAdapter) in info.parametersAdapters.withIndex()) {
                    jsClosureCallerCall.putValueArgument(
                        adapterIndex + 1,
                        paramAdapter.adaptIfNeeded(
                            irGet(valueParameters[adapterIndex]),
                            lambdaBuilder
                        )
                    )
                }
                +irReturn(info.resultAdapter.adaptIfNeeded(jsClosureCallerCall, lambdaBuilder))
            }

            overriddenSymbols =
                overriddenSymbols + functionType.classOrNull!!.functions.single { it.owner.name == Name.identifier("invoke") }
        }

        val builder = context.createIrBuilder(result.symbol)
        result.body = builder.irBlockBody {
            +irReturn(irCall(closureClassConstructor).also { it.putValueArgument(0, irGet(result.valueParameters[0])) })
        }

        additionalDeclarations += closureClass
        additionalDeclarations += result
        return result
    }

    private fun createJsClosureCaller(info: FunctionTypeInfo): IrSimpleFunction {
        val result = context.irFactory.buildFun {
            name = Name.identifier("__callJsClosure_${info.hashString}")
            returnType = info.adaptedResultType
        }
        result.parent = currentParent
        result.addValueParameter {
            name = Name.identifier("f")
            type = symbols.externalInterfaceType
        }
        val arity = info.adaptedParameterTypes.size
        repeat(arity) { paramIndex ->
            result.addValueParameter {
                name = Name.identifier("p$paramIndex")
                type = info.adaptedParameterTypes[paramIndex]
            }
        }
        val builder = context.createIrBuilder(result.symbol)
        val jsFun = buildString {
            append("(f, ")
            appendParameterList(arity)
            append(") => f(")
            appendParameterList(arity)
            append(")")
        }

        result.annotations += builder.irCallConstructor(context.wasmSymbols.jsFunConstructor, typeArguments = emptyList()).also {
            it.putValueArgument(0, builder.irString(jsFun))
        }

        additionalDeclarations += result
        return result
    }

    inner class FunctionTypeInfo(val functionType: IrSimpleType, toJs: Boolean) {
        init {
            require(functionType.arguments.all { it is IrTypeProjection }) {
                "Star projection is not supported in function type interop ${functionType.render()}"
            }
        }

        val hashString: String =
            functionType.hashCode().absoluteValue.toString(Character.MAX_RADIX)

        val originalParameterTypes: List<IrType> =
            functionType.arguments.dropLast(1).map { (it as IrTypeProjection).type }

        val originalResultType: IrType =
            (functionType.arguments.last() as IrTypeProjection).type

        val parametersAdapters: List<InteropTypeAdapter?> =
            originalParameterTypes.map { parameterType ->
                if (toJs)
                    parameterType.jsToKotlinAdapterIfNeeded(isReturn = false)
                else
                    parameterType.kotlinToJsAdapterIfNeeded(isReturn = false)
            }

        val resultAdapter: InteropTypeAdapter? =
            if (toJs)
                originalResultType.kotlinToJsAdapterIfNeeded(isReturn = true)
            else
                originalResultType.jsToKotlinAdapterIfNeeded(isReturn = true)

        val adaptedParameterTypes: List<IrType> =
            originalParameterTypes.zip(parametersAdapters).map { (parameterType, adapter) ->
                (if (toJs) adapter?.fromType else adapter?.toType) ?: parameterType
            }

        val adaptedResultType: IrType =
            (if (toJs) resultAdapter?.toType else resultAdapter?.fromType) ?: originalResultType
    }

    interface InteropTypeAdapter {
        val fromType: IrType
        val toType: IrType
        fun adapt(expression: IrExpression, builder: IrBuilderWithScope): IrExpression
    }

    fun InteropTypeAdapter?.adaptIfNeeded(expression: IrExpression, builder: IrBuilderWithScope): IrExpression =
        this?.adapt(expression, builder) ?: expression

    /**
     * Adapter implemented as a single function call
     */
    class FunctionBasedAdapter(
        private val function: IrSimpleFunction,
    ) : InteropTypeAdapter {
        override val fromType = function.valueParameters[0].type
        override val toType = function.returnType
        override fun adapt(expression: IrExpression, builder: IrBuilderWithScope): IrExpression {
            val call = builder.irCall(function)
            call.putValueArgument(0, expression)
            return call
        }
    }

    /**
     * Current V8 Wasm GC mandates dataref type instead of structs and arrays
     */
    inner class SendKotlinObjectToJsAdapter(
        override val fromType: IrType
    ) : InteropTypeAdapter {
        override val toType: IrType = context.wasmSymbols.wasmDataRefType
        override fun adapt(expression: IrExpression, builder: IrBuilderWithScope): IrExpression {
            return builder.irReinterpretCast(expression, toType)
        }
    }

    /**
     * Current V8 Wasm GC mandates dataref type instead of structs and arrays
     */
    inner class ReceivingKotlinObjectFromJsAdapter(
        override val toType: IrType
    ) : InteropTypeAdapter {
        override val fromType: IrType = context.wasmSymbols.wasmDataRefType
        override fun adapt(expression: IrExpression, builder: IrBuilderWithScope): IrExpression {
            val call = builder.irCall(context.wasmSymbols.refCast)
            call.putValueArgument(0, expression)
            call.putTypeArgument(0, toType)
            return call
        }
    }

    /**
     * Effectively `value?.let { adapter(it) }`
     */
    inner class NullOrAdapter(
        val adapter: InteropTypeAdapter
    ) : InteropTypeAdapter {
        override val fromType: IrType = adapter.fromType.makeNullable()
        override val toType: IrType = adapter.toType.makeNullable()
        override fun adapt(expression: IrExpression, builder: IrBuilderWithScope): IrExpression {
            return builder.irComposite {
                val tmp = irTemporary(adapter.adapt(expression, builder))
                +irIfNull(toType, irGet(tmp), irNull(toType), irImplicitCast(irGet(tmp), toType))
            }
        }
    }
}

internal fun StringBuilder.appendParameterList(size: Int, name: String = "p", isEnd: Boolean = true) =
    repeat(size) {
        append(name)
        append(it)
        if (!isEnd || it + 1 < size)
            append(", ")
    }

/**
 * Redirect calls to external and @JsExport functions to created wrappers
 */
class JsInteropFunctionCallsLowering(val context: WasmBackendContext) : BodyLoweringPass {
    override fun lower(irBody: IrBody, container: IrDeclaration) {
        irBody.transformChildrenVoid(object : IrElementTransformerVoid() {
            override fun visitCall(expression: IrCall): IrExpression {
                expression.transformChildrenVoid()
                val newFun: IrSimpleFunction? = context.mapping.wasmJsInteropFunctionToWrapper[expression.symbol.owner]
                return if (newFun != null && container != newFun) {
                    irCall(expression, newFun)
                } else {
                    expression
                }
            }
        })
    }
}