/*
 * Copyright 2010-2025 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.common.lower

import org.jetbrains.kotlin.backend.common.BodyLoweringPass
import org.jetbrains.kotlin.backend.common.CommonBackendContext
import org.jetbrains.kotlin.backend.common.functionWithContinuations
import org.jetbrains.kotlin.descriptors.DescriptorVisibilities
import org.jetbrains.kotlin.ir.IrStatement
import org.jetbrains.kotlin.ir.UNDEFINED_OFFSET
import org.jetbrains.kotlin.ir.builders.IrBuilderWithScope
import org.jetbrains.kotlin.ir.builders.declarations.addConstructor
import org.jetbrains.kotlin.ir.builders.declarations.addField
import org.jetbrains.kotlin.ir.builders.declarations.addFunction
import org.jetbrains.kotlin.ir.builders.declarations.buildClass
import org.jetbrains.kotlin.ir.builders.declarations.buildValueParameter
import org.jetbrains.kotlin.ir.builders.irBlock
import org.jetbrains.kotlin.ir.builders.irBlockBody
import org.jetbrains.kotlin.ir.builders.irCallConstructor
import org.jetbrains.kotlin.ir.builders.irExprBody
import org.jetbrains.kotlin.ir.builders.irGet
import org.jetbrains.kotlin.ir.builders.irGetField
import org.jetbrains.kotlin.ir.builders.irReturn
import org.jetbrains.kotlin.ir.builders.irTemporary
import org.jetbrains.kotlin.ir.builders.setSourceRange
import org.jetbrains.kotlin.ir.declarations.IrClass
import org.jetbrains.kotlin.ir.declarations.IrConstructor
import org.jetbrains.kotlin.ir.declarations.IrDeclaration
import org.jetbrains.kotlin.ir.declarations.IrDeclarationBase
import org.jetbrains.kotlin.ir.declarations.IrDeclarationOrigin
import org.jetbrains.kotlin.ir.declarations.IrDeclarationParent
import org.jetbrains.kotlin.ir.declarations.IrField
import org.jetbrains.kotlin.ir.declarations.IrParameterKind
import org.jetbrains.kotlin.ir.declarations.IrSimpleFunction
import org.jetbrains.kotlin.ir.declarations.IrSymbolOwner
import org.jetbrains.kotlin.ir.declarations.IrValueParameter
import org.jetbrains.kotlin.ir.expressions.IrBlockBody
import org.jetbrains.kotlin.ir.expressions.IrBody
import org.jetbrains.kotlin.ir.expressions.IrDelegatingConstructorCall
import org.jetbrains.kotlin.ir.expressions.IrExpression
import org.jetbrains.kotlin.ir.expressions.IrExpressionBody
import org.jetbrains.kotlin.ir.expressions.IrFunctionReference
import org.jetbrains.kotlin.ir.expressions.IrReturn
import org.jetbrains.kotlin.ir.expressions.IrRichFunctionReference
import org.jetbrains.kotlin.ir.expressions.IrStatementOrigin
import org.jetbrains.kotlin.ir.expressions.impl.IrInstanceInitializerCallImpl
import org.jetbrains.kotlin.ir.expressions.implicitCastTo
import org.jetbrains.kotlin.ir.symbols.IrClassSymbol
import org.jetbrains.kotlin.ir.types.IrSimpleType
import org.jetbrains.kotlin.ir.types.IrType
import org.jetbrains.kotlin.ir.types.IrTypeProjection
import org.jetbrains.kotlin.ir.types.IrTypeSubstitutor
import org.jetbrains.kotlin.ir.types.classOrFail
import org.jetbrains.kotlin.ir.types.extractTypeParameters
import org.jetbrains.kotlin.ir.types.isNothing
import org.jetbrains.kotlin.ir.types.typeWith
import org.jetbrains.kotlin.ir.util.addFakeOverrides
import org.jetbrains.kotlin.ir.util.copyTo
import org.jetbrains.kotlin.ir.util.createDispatchReceiverParameterWithClassParent
import org.jetbrains.kotlin.ir.util.createThisReceiverParameter
import org.jetbrains.kotlin.ir.util.defaultType
import org.jetbrains.kotlin.ir.util.erasedUpperBound
import org.jetbrains.kotlin.ir.util.isLambda
import org.jetbrains.kotlin.ir.util.nonDispatchParameters
import org.jetbrains.kotlin.ir.util.primaryConstructor
import org.jetbrains.kotlin.ir.visitors.IrTransformer
import org.jetbrains.kotlin.name.Name
import org.jetbrains.kotlin.types.Variance
import org.jetbrains.kotlin.utils.addToStdlib.applyIf
import org.jetbrains.kotlin.utils.addToStdlib.shouldNotBeCalled
import org.jetbrains.kotlin.utils.memoryOptimizedPlus
import kotlin.collections.plus

/**
 * This lowering transforms [IrRichFunctionReference] nodes to an anonymous class.
 *
 * The class would have:
 *   * Constructor capturing all values from [IrRichFunctionReference.boundValues], and storing them to fields
 *   * A method overriding [IrRichFunctionReference.overriddenFunctionSymbol], with body moved from [IrRichFunctionReference.invokeFunction]
 *   * [IrRichFunctionReference.type] as a super-interface type (typically `[K][Suspend]FunctionN`,
 *     or fun interface the reference was sam converted to)
 *
 * Platforms can customize:
 *   * Super-class with platform-specific reference implementation details by overriding [getSuperClassType] method
 *   * Equality/hashCode/toString and other reflection methods implementation by adding methods in overridden [generateExtraMethods] method
 *   * exact names/origins of generated classes/methods by overriding corresponding methods
 *
 * For example, the following code:
 * ```kotlin
 * fun foo1(l: () -> String): String {
 *   return l()
 * }
 * fun <FooTP> foo2(v: FooTP, l: (FooTP) -> String): String {
 *   return l(v)
 * }
 *
 * private fun <T> bar(t: T): String { /* ... */ }
 *
 * fun <BarTP> bar(v: BarTP): String {
 *   return foo1(v::bar/*<T=BarTP>*/) + foo2(v) { bar/*<T=BarTP>*/(it) }
 * }
 * ```
 *
 * is lowered into:
 * ```kotlin
 * fun <BarTP> bar(v: BarTP): String {
 *   /*local*/ class <local-platform-specific-name-1>(p$0: BarTP) : KFunction0<String>, PlatformSpecificSuperType() {
 *     private val f$0: BarTP = p$0
 *     override fun invoke() = bar<BarTP>(f$0)
 *     // some platform specific reflection information
 *   }
 *   /*local*/ class <local-platform-specific-name-2> : Function1<BarTP, String>, PlatformSpecificSuperTypeProbablyAny() {
 *     override fun invoke(p0: BarTP) = bar<BarTP>(p0)
 *   }
 *   return foo1(<local-platform-specific-name-1>(v)) + foo2(v, <local-platform-specific-name-2>())
 * }
 * ```
 *
 * Note that as all these classes are defined as local ones, they don't need to explicitly capture local variables or any type parameters.
 * But it can happen, that [LocalDeclarationsLowering] would later capture something additional into the classes.
 */
abstract class AbstractFunctionReferenceLowering<C : CommonBackendContext>(val context: C) : BodyLoweringPass {
    override fun lower(irBody: IrBody, container: IrDeclaration) {
        irBody.transform(object : IrTransformer<IrDeclarationParent>() {
            override fun visitDeclaration(declaration: IrDeclarationBase, data: IrDeclarationParent): IrStatement {
                return super.visitDeclaration(declaration, declaration as? IrDeclarationParent ?: data)
            }

            override fun visitRichFunctionReference(expression: IrRichFunctionReference, data: IrDeclarationParent): IrExpression {
                expression.transformChildren(this, data)
                val irBuilder = context.createIrBuilder(
                    (data as IrSymbolOwner).symbol,
                    expression.startOffset, expression.endOffset
                )

                val clazz = buildClass(expression, data)
                val constructor = clazz.primaryConstructor!!
                val newExpression = irBuilder.irCallConstructor(constructor.symbol, emptyList()).apply {
                    origin = getConstructorCallOrigin(expression)
                    for ((index, value) in expression.boundValues.withIndex()) {
                        arguments[index] = value
                    }
                    for (index in expression.boundValues.size until arguments.size) {
                        arguments[index] = irBuilder.getExtraConstructorArgument(constructor.parameters[index], expression)
                    }
                }
                return irBuilder.irBlock {
                    +clazz
                    +newExpression
                }
            }

            override fun visitFunctionReference(expression: IrFunctionReference, data: IrDeclarationParent): IrExpression {
                shouldNotBeCalled()
            }
        }, container as? IrDeclarationParent ?: container.parent)
    }

    // SAM class used as a superclass can sometimes have type projections.
    // But that's not suitable for super-types, so we erase them
    protected fun IrType.removeProjections(): IrType {
        if (this !is IrSimpleType) return this
        val arguments = arguments.mapIndexed { index, argument ->
            val typeParameter = (classifier as IrClassSymbol).owner.typeParameters[index]
            fun erasedUpperBound() = typeParameter.erasedUpperBound.defaultType

            // Star projections are not allowed in supertype clause
            if (argument !is IrTypeProjection) return@mapIndexed erasedUpperBound()

            // `in` and `out` projections are not allowed either
            if (argument.variance != Variance.INVARIANT) return@mapIndexed erasedUpperBound()

            // In case a lambda parameter's type is inferred by the frontend to an intersection type, in IR
            // it will be approximated to `Nothing` (because intersection types are not representable in IR at all).
            // Since function parameters are contravariant, `Nothing` is the only type we can approximate an intersection type to.
            // We cannot use `Nothing` as a parameter type, though —
            // semantically it would mean that such a function reference can never be invoked, which is not true.
            // Some targets like Wasm can break because of this.
            // That's why we treat such type arguments similarly to type projections.
            // For a concrete example, see this test: compiler/testData/codegen/box/callableReference/kt49526_sam.kt
            if (typeParameter.variance == Variance.IN_VARIANCE && argument.type.isNothing()) return@mapIndexed erasedUpperBound()

            argument.type
        }
        return classifier.typeWith(arguments)
    }

    private fun buildClass(functionReference: IrRichFunctionReference, parent: IrDeclarationParent): IrClass {
        val functionReferenceClass = context.irFactory.buildClass {
            startOffset = functionReference.startOffset
            endOffset = functionReference.endOffset
            origin = getClassOrigin(functionReference)
            name = getReferenceClassName(functionReference)
            visibility = DescriptorVisibilities.LOCAL
        }.apply {
            this.parent = parent
            createThisReceiverParameter()
        }
        val superClass = getSuperClassType(functionReference)
        val superInterfaceType = functionReference.type.removeProjections()
        functionReferenceClass.superTypes =
            listOf(superClass, superInterfaceType) memoryOptimizedPlus getAdditionalInterfaces(functionReference)
        val constructor = functionReferenceClass.addConstructor {
            origin = getConstructorOrigin(functionReference)
            isPrimary = true
        }.apply {
            parameters = functionReference.boundValues.mapIndexed { index, value ->
                buildValueParameter(this) {
                    name = Name.identifier("p${index}")
                    startOffset = value.startOffset
                    endOffset = value.endOffset
                    type = value.type
                    kind = IrParameterKind.Regular
                }
            } + getExtraConstructorParameters(this, functionReference)
            body = context.createIrBuilder(symbol, this.startOffset, this.endOffset).irBlockBody {
                +generateSuperClassConstructorCall(this@apply, superClass, functionReference)
                +IrInstanceInitializerCallImpl(this.startOffset, this.endOffset, functionReferenceClass.symbol, context.irBuiltIns.unitType)
            }
        }

        val fields = functionReference.boundValues.mapIndexed { index, captured ->
            functionReferenceClass.addField {
                startOffset = captured.startOffset
                endOffset = captured.endOffset
                name = Name.identifier("f\$${index}")
                visibility = DescriptorVisibilities.PRIVATE
                isFinal = true
                type = captured.type
            }.apply {
                val builder = context.createIrBuilder(symbol, startOffset, endOffset)
                initializer = builder.irExprBody(builder.irGet(constructor.parameters[index]))
            }
        }
        buildInvokeMethod(
            functionReference,
            functionReferenceClass,
            superInterfaceType,
            fields,
        ).apply {
            postprocessInvoke(this, functionReference)
        }

        generateExtraMethods(functionReferenceClass, functionReference)

        functionReferenceClass.addFakeOverrides(
            context.typeSystem,
            // Built function overrides originalSuperMethod, while, if parent class is already lowered, it would
            // transformedSuperMethod in its declaration list. We need not fake override in that case.
            // Later lowerings will fix it and replace function with one overriding transformedSuperMethod.
            ignoredParentSymbols = listOfNotNull(functionReference.overriddenFunctionSymbol.owner.functionWithContinuations?.symbol),
        )
        postprocessClass(functionReferenceClass, functionReference)
        return functionReferenceClass
    }

    private fun buildInvokeMethod(
        functionReference: IrRichFunctionReference,
        functionReferenceClass: IrClass,
        superInterfaceType: IrType,
        boundFields: List<IrField>
    ): IrSimpleFunction {
        val superFunction = functionReference.overriddenFunctionSymbol.owner
        val invokeFunction = functionReference.invokeFunction
        val isLambda = functionReference.origin.isLambda
        return functionReferenceClass.addFunction {
            setSourceRange(if (isLambda) invokeFunction else functionReference)
            origin = getInvokeMethodOrigin(functionReference)
            name = superFunction.name
            returnType = invokeFunction.returnType
            isOperator = superFunction.isOperator
            isSuspend = superFunction.isSuspend
        }.apply {
            attributeOwnerId = functionReference.attributeOwnerId
            annotations = invokeFunction.annotations

            parameters += createDispatchReceiverParameterWithClassParent()
            require(superFunction.typeParameters.isEmpty()) { "Fun interface abstract function can't have type parameters" }

            val typeSubstitutor = IrTypeSubstitutor(
                extractTypeParameters(superInterfaceType.classOrFail.owner).map { it.symbol },
                (superInterfaceType as IrSimpleType).arguments,
                allowEmptySubstitution = true
            )

            val nonDispatchParameters = superFunction.nonDispatchParameters.mapIndexed { i, superParameter ->
                superParameter.copyTo(
                    this,
                    startOffset = if (isLambda) invokeFunction.parameters[i].startOffset else UNDEFINED_OFFSET,
                    endOffset = if (isLambda) invokeFunction.parameters[i].endOffset else UNDEFINED_OFFSET,
                    name = invokeFunction.parameters[i].name,
                    type = typeSubstitutor.substitute(superParameter.type),
                    defaultValue = null,
                )
            }
            this.parameters += nonDispatchParameters
            overriddenSymbols += superFunction.symbol

            val builder = context.createIrBuilder(symbol).applyIf(isLambda) { at(invokeFunction.body!!) }
            body = builder.irBlockBody {
                val variablesMapping = buildMap {
                    for ((index, field) in boundFields.withIndex()) {
                        put(invokeFunction.parameters[index], irTemporary(irGetField(irGet(dispatchReceiverParameter!!), field)))
                    }
                    for ((index, parameter) in nonDispatchParameters.withIndex()) {
                        val invokeParameter = invokeFunction.parameters[index + boundFields.size]
                        if (parameter.type != invokeParameter.type) {
                            put(invokeParameter, irTemporary(irGet(parameter).implicitCastTo(invokeParameter.type)))
                        } else {
                            put(invokeParameter, parameter)
                        }
                    }
                }
                val transformedBody = invokeFunction.body!!.transform(object : VariableRemapper(variablesMapping) {
                    override fun visitReturn(expression: IrReturn): IrExpression {
                        if (expression.returnTargetSymbol == invokeFunction.symbol) {
                            expression.returnTargetSymbol = this@apply.symbol
                        }
                        return super.visitReturn(expression)
                    }

                    override fun visitDeclaration(declaration: IrDeclarationBase): IrStatement {
                        if (declaration.parent == invokeFunction)
                            declaration.parent = this@apply
                        return super.visitDeclaration(declaration)
                    }
                }, null)
                when (transformedBody) {
                    is IrBlockBody -> +transformedBody.statements
                    is IrExpressionBody -> +irReturn(transformedBody.expression)
                    else -> error("Unexpected body type: ${transformedBody::class.simpleName}")
                }
            }
        }
    }

    protected open fun postprocessClass(functionReferenceClass: IrClass, functionReference: IrRichFunctionReference) {}
    protected open fun postprocessInvoke(invokeFunction: IrSimpleFunction, functionReference: IrRichFunctionReference) {}
    protected open fun generateExtraMethods(functionReferenceClass: IrClass, reference: IrRichFunctionReference) {}

    protected open fun getExtraConstructorParameters(
        constructor: IrConstructor,
        reference: IrRichFunctionReference,
    ): List<IrValueParameter> = emptyList()

    protected open fun IrBuilderWithScope.getExtraConstructorArgument(
        parameter: IrValueParameter,
        reference: IrRichFunctionReference,
    ): IrExpression? = null

    protected abstract fun IrBuilderWithScope.generateSuperClassConstructorCall(
        constructor: IrConstructor,
        superClassType: IrType,
        functionReference: IrRichFunctionReference,
    ): IrDelegatingConstructorCall

    protected abstract fun getReferenceClassName(reference: IrRichFunctionReference): Name
    protected abstract fun getSuperClassType(reference: IrRichFunctionReference): IrType
    protected open fun getAdditionalInterfaces(reference: IrRichFunctionReference): List<IrType> = emptyList()
    protected abstract fun getClassOrigin(reference: IrRichFunctionReference): IrDeclarationOrigin
    protected abstract fun getConstructorOrigin(reference: IrRichFunctionReference): IrDeclarationOrigin
    protected abstract fun getInvokeMethodOrigin(reference: IrRichFunctionReference): IrDeclarationOrigin
    protected abstract fun getConstructorCallOrigin(reference: IrRichFunctionReference): IrStatementOrigin?
}
