/*
 * Copyright 2023 The Android Open Source Project
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package androidx.compose.compiler.plugins.kotlin.lower.hiddenfromobjc

import androidx.compose.compiler.plugins.kotlin.FeatureFlags
import androidx.compose.compiler.plugins.kotlin.ModuleMetrics
import androidx.compose.compiler.plugins.kotlin.analysis.StabilityInferencer
import androidx.compose.compiler.plugins.kotlin.lower.AbstractComposeLowering
import androidx.compose.compiler.plugins.kotlin.lower.containsComposableAnnotation
import androidx.compose.compiler.plugins.kotlin.lower.hasFirDeclaration
import androidx.compose.compiler.plugins.kotlin.lower.needsComposableRemapping
import org.jetbrains.kotlin.backend.common.extensions.IrPluginContext
import org.jetbrains.kotlin.descriptors.DescriptorVisibilities
import org.jetbrains.kotlin.ir.IrStatement
import org.jetbrains.kotlin.ir.declarations.*
import org.jetbrains.kotlin.ir.expressions.impl.IrConstructorCallImpl
import org.jetbrains.kotlin.ir.expressions.impl.fromSymbolOwner
import org.jetbrains.kotlin.ir.symbols.IrClassSymbol
import org.jetbrains.kotlin.ir.types.defaultType
import org.jetbrains.kotlin.ir.util.constructors
import org.jetbrains.kotlin.ir.util.isLocal
import org.jetbrains.kotlin.ir.visitors.transformChildrenVoid
import org.jetbrains.kotlin.name.ClassId
import org.jetbrains.kotlin.platform.konan.isNative

val hiddenFromObjCClassId = ClassId.fromString("kotlin/native/HiddenFromObjC")

/**
 *  AddHiddenFromObjCLowering looks for functions and properties with @Composable types and
 *  adds the `kotlin.native.HiddenFromObjC` annotation to them.
 *  [docs](https://kotlinlang.org/api/latest/jvm/stdlib/kotlin.native/-hidden-from-obj-c/)
 */
class AddHiddenFromObjCLowering(
    private val pluginContext: IrPluginContext,
    metrics: ModuleMetrics,
    private val hideFromObjCDeclarationsSet: HideFromObjCDeclarationsSet?,
    stabilityInferencer: StabilityInferencer,
    featureFlags: FeatureFlags,
) : AbstractComposeLowering(
    pluginContext,
    metrics,
    stabilityInferencer,
    featureFlags
) {

    private val hiddenFromObjCAnnotation: IrClassSymbol by lazy {
        getTopLevelClass(hiddenFromObjCClassId)
    }

    private var currentShouldAnnotateClass = false

    override fun lower(irModule: IrModuleFragment) {
        require(context.platform.isNative()) {
            "AddHiddenFromObjCLowering is expected to run only for k/native. " +
                    "The platform - ${context.platform}"
        }
        irModule.transformChildrenVoid(this)
    }

    /** `visitClass` is only needed until [issue](https://youtrack.jetbrains.com/issue/KT-65288/) fix
     * after the issue is resolved, `visitClass` could be removed entirely
     */
    override fun visitClass(declaration: IrClass): IrStatement {
        val previousShouldAnnotateClass = currentShouldAnnotateClass
        currentShouldAnnotateClass = false

        val cls = super.visitClass(declaration) as IrClass

        // We see an issue only with data classes containing something Composable.
        // Adding an annotation to all classes makes the FirNativeHiddenFromObjCInheritanceChecker (kotlin) complain.
        // data classes can't be open, so it should work.
        if (currentShouldAnnotateClass && cls.isData) {
            cls.addHiddenFromObjCAnnotation()
            hideFromObjCDeclarationsSet?.add(cls)
        }

        currentShouldAnnotateClass = previousShouldAnnotateClass
        return cls
    }

    private fun IrFunction.isSyntheticFun(): Boolean =
        origin == IrDeclarationOrigin.FAKE_OVERRIDE || startOffset < 0 || endOffset < 0


    override fun visitFunction(declaration: IrFunction): IrStatement {
        val f = super.visitFunction(declaration) as IrFunction
        if (f.isLocal || f.isSyntheticFun() ||
            !(f.visibility == DescriptorVisibilities.PUBLIC ||
                    f.visibility == DescriptorVisibilities.PROTECTED)
        )
            return f

        if (f.hasComposableAnnotation() || f.needsComposableRemapping()) {
            f.addHiddenFromObjCAnnotation()
            hideFromObjCDeclarationsSet?.add(f)
            currentShouldAnnotateClass = true
        }

        return f
    }

    override fun visitProperty(declaration: IrProperty): IrStatement {
        val p = super.visitProperty(declaration) as IrProperty
        if (p.isLocal || p.getter?.isSyntheticFun() == true || p.visibility != DescriptorVisibilities.PUBLIC) return p

        val shouldAdd = p.getter?.hasComposableAnnotation() ?: false ||
                p.getter?.needsComposableRemapping() ?: false ||
                p.backingField?.type.containsComposableAnnotation()

        if (shouldAdd) {
            p.addHiddenFromObjCAnnotation()
            hideFromObjCDeclarationsSet?.add(p)
            currentShouldAnnotateClass = true
        }

        return p
    }

    private fun IrDeclaration.addHiddenFromObjCAnnotation() {
        if (!hasFirDeclaration()) {
            return
        }
        val annotation = IrConstructorCallImpl.fromSymbolOwner(
            type = hiddenFromObjCAnnotation.defaultType,
            constructorSymbol = hiddenFromObjCAnnotation.constructors.first()
        )
        pluginContext.metadataDeclarationRegistrar.addMetadataVisibleAnnotationsToElement(this, annotation)
    }
}
