/*
 * Copyright 2010-2024 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.analysis.api.fir.components

import org.jetbrains.kotlin.analysis.api.components.KaCompilerPluginGeneratedDeclarations
import org.jetbrains.kotlin.analysis.api.components.KaCompilerPluginGeneratedDeclarationsProvider
import org.jetbrains.kotlin.analysis.api.fir.KaFirSession
import org.jetbrains.kotlin.analysis.api.fir.KaSymbolByFirBuilder
import org.jetbrains.kotlin.analysis.api.impl.base.components.KaBaseCompilerPluginGeneratedDeclarations
import org.jetbrains.kotlin.analysis.api.impl.base.components.KaBaseSessionComponent
import org.jetbrains.kotlin.analysis.api.impl.base.scopes.KaBaseEmptyScope
import org.jetbrains.kotlin.analysis.api.lifetime.KaLifetimeToken
import org.jetbrains.kotlin.analysis.api.lifetime.withValidityAssertion
import org.jetbrains.kotlin.analysis.api.projectStructure.KaModule
import org.jetbrains.kotlin.analysis.api.scopes.KaScope
import org.jetbrains.kotlin.analysis.api.symbols.KaCallableSymbol
import org.jetbrains.kotlin.analysis.api.symbols.KaClassifierSymbol
import org.jetbrains.kotlin.analysis.api.symbols.KaConstructorSymbol
import org.jetbrains.kotlin.analysis.api.symbols.KaPackageSymbol
import org.jetbrains.kotlin.fir.extensions.FirSwitchableExtensionDeclarationsSymbolProvider
import org.jetbrains.kotlin.fir.extensions.generatedDeclarationsSymbolProvider
import org.jetbrains.kotlin.fir.resolve.providers.FirSymbolNamesProvider
import org.jetbrains.kotlin.name.ClassId
import org.jetbrains.kotlin.name.FqName
import org.jetbrains.kotlin.name.Name
import kotlin.collections.orEmpty

internal class KaFirCompilerPluginGeneratedDeclarationsProvider(
    override val analysisSessionProvider: () -> KaFirSession,
) : KaBaseSessionComponent<KaFirSession>(), KaCompilerPluginGeneratedDeclarationsProvider {
    override val KaModule.compilerPluginGeneratedDeclarations: KaCompilerPluginGeneratedDeclarations
        get() = withValidityAssertion {
            val firSessionForModule = analysisSession.resolutionFacade.sessionProvider.getSession(this)
            val generatedDeclarationsSymbolProviderForModule = firSessionForModule.generatedDeclarationsSymbolProvider
                ?: return KaBaseCompilerPluginGeneratedDeclarations(KaBaseEmptyScope(analysisSession.token))

            val topLevelScope = KaFirTopLevelCompilerPluginGeneratedDeclarationsScope(
                generatedDeclarationsSymbolProviderForModule,
                analysisSession.firSymbolBuilder,
            )

            KaBaseCompilerPluginGeneratedDeclarations(topLevelScope)
        }
}

/**
 * A [KaScope] implementation containing all top-level declarations generated by compiler plugins.
 *
 * Relies on [generatedDeclarationsSymbolProvider] to retrieve the generated declarations.
 */
private class KaFirTopLevelCompilerPluginGeneratedDeclarationsScope(
    private val generatedDeclarationsSymbolProvider: FirSwitchableExtensionDeclarationsSymbolProvider,
    private val symbolByFirBuilder: KaSymbolByFirBuilder,
) : KaScope {

    override val token: KaLifetimeToken get() = symbolByFirBuilder.token

    /**
     * N.B. We expect that this [FirSymbolNamesProvider] is aware of all the names
     * potentially generated by the compiler plugins.
     * Otherwise, the whole scope implementation would be incorrect.
     */
    private val symbolNamesProvider: FirSymbolNamesProvider = generatedDeclarationsSymbolProvider.symbolNamesProvider

    override fun callables(nameFilter: (Name) -> Boolean): Sequence<KaCallableSymbol> = withValidityAssertion {
        callablesImpl { packageFqName ->
            symbolNamesProvider.getTopLevelCallableNamesInPackage(packageFqName)?.filter(nameFilter).orEmpty()
        }
    }

    override fun callables(names: Collection<Name>): Sequence<KaCallableSymbol> = withValidityAssertion {
        if (names.isEmpty()) return emptySequence()

        callablesImpl { _ -> names }
    }

    private fun callablesImpl(
        callableNamesForPackage: (packageFqName: FqName) -> Collection<Name>,
    ): Sequence<KaCallableSymbol> = sequence {
        val packagesWithCallables = symbolNamesProvider.getPackageNamesWithTopLevelCallables().orEmpty()

        for (packageName in packagesWithCallables) {
            val packageFqName = FqName(packageName)

            val callableNames = callableNamesForPackage(packageFqName)

            for (callableName in callableNames) {
                val callables = generatedDeclarationsSymbolProvider.getTopLevelCallableSymbols(packageFqName, callableName)

                for (callable in callables) {
                    val callableSymbol = symbolByFirBuilder.callableBuilder.buildCallableSymbol(callable)

                    yield(callableSymbol)
                }
            }
        }
    }

    override fun classifiers(nameFilter: (Name) -> Boolean): Sequence<KaClassifierSymbol> = withValidityAssertion {
        classifiersImpl { packageFqName ->
            symbolNamesProvider.getTopLevelClassifierNamesInPackage(packageFqName)?.filter(nameFilter).orEmpty()
        }
    }

    override fun classifiers(names: Collection<Name>): Sequence<KaClassifierSymbol> = withValidityAssertion {
        if (names.isEmpty()) return emptySequence()

        classifiersImpl { _ -> names }
    }

    private fun classifiersImpl(
        classifierNamesForPackage: (FqName) -> Collection<Name>,
    ): Sequence<KaClassifierSymbol> = sequence {
        val packagesWithClassifiers = symbolNamesProvider.getPackageNamesWithTopLevelClassifiers().orEmpty()

        for (packageName in packagesWithClassifiers) {
            val packageFqName = FqName(packageName)

            val classifierNames = classifierNamesForPackage(packageFqName)

            for (classifierName in classifierNames) {
                val classifierId = ClassId(packageFqName, classifierName)
                val classifier = generatedDeclarationsSymbolProvider.getClassLikeSymbolByClassId(classifierId) ?: continue
                val classifierSymbol = symbolByFirBuilder.classifierBuilder.buildClassLikeSymbol(classifier)

                yield(classifierSymbol)
            }
        }
    }

    override val constructors: Sequence<KaConstructorSymbol>
        get() = withValidityAssertion {
            // constructors are not considered to be top-level declarations
            emptySequence()
        }

    override fun getPackageSymbols(nameFilter: (Name) -> Boolean): Sequence<KaPackageSymbol> = withValidityAssertion {
        // we do not provide generated packages at all, since the semantics and use-cases are not yet clear
        emptySequence()
    }

    override fun getPossibleCallableNames(): Set<Name> = withValidityAssertion {
        buildSet {
            val packagesWithCallables = symbolNamesProvider.getPackageNamesWithTopLevelCallables().orEmpty()

            for (packageName in packagesWithCallables) {
                val packageFqName = FqName(packageName)
                val callableNames = symbolNamesProvider.getTopLevelCallableNamesInPackage(packageFqName) ?: continue

                addAll(callableNames)
            }
        }
    }

    override fun getPossibleClassifierNames(): Set<Name> = withValidityAssertion {
        buildSet {
            val packagesWithClassifiers = symbolNamesProvider.getPackageNamesWithTopLevelClassifiers().orEmpty()

            for (packageName in packagesWithClassifiers) {
                val packageFqName = FqName(packageName)
                val classifierNames = symbolNamesProvider.getTopLevelClassifierNamesInPackage(packageFqName) ?: continue

                addAll(classifierNames)
            }
        }
    }
}