/*
 * 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.kotlin.analysis.api.fir.scopes

import com.intellij.psi.PsiManager
import com.intellij.psi.impl.file.PsiPackageImpl
import org.jetbrains.kotlin.analysis.api.fir.KaFirSession
import org.jetbrains.kotlin.analysis.api.platform.utils.forEachNonKotlinPsiElementFinder
import org.jetbrains.kotlin.fir.FirSession
import org.jetbrains.kotlin.fir.extensions.ExperimentalTopLevelDeclarationsGenerationApi
import org.jetbrains.kotlin.fir.extensions.declarationGenerators
import org.jetbrains.kotlin.fir.extensions.extensionService
import org.jetbrains.kotlin.fir.resolve.providers.FirSymbolNamesProvider
import org.jetbrains.kotlin.fir.resolve.providers.symbolProvider
import org.jetbrains.kotlin.name.FqName
import org.jetbrains.kotlin.name.Name
import org.jetbrains.kotlin.platform.jvm.isJvm

/**
 * Provides top-level names for classifiers and callables in given packages. Apart from names found in sources and binaries,
 * [DeclarationsInPackageProvider] also collects names for classifiers and callables generated by
 * [FirDeclarationGenerationExtension][org.jetbrains.kotlin.fir.extensions.FirDeclarationGenerationExtension]s.
 *
 * [DeclarationsInPackageProvider] first tries the [FirSymbolNamesProvider] of the analysis session's underlying FIR session, because it
 * provides symbol names from binary libraries in Standalone mode, which the Standalone declaration provider does not contain (to avoid
 * stub-indexing binary libraries). And in general, querying the symbol names provider might be faster since its sets are cached, which is
 * not necessarily the case for declaration providers (e.g. the IDE declaration provider hitting the index without caching).
 *
 * However, since symbol names providers may not be able to compute name sets per their contract, we may have to fall back to the
 * [declaration provider][org.jetbrains.kotlin.analysis.api.platform.declarations.KotlinDeclarationProvider].
 *
 * Since this fallback is not suitable for symbols from binary libraries in Standalone mode, the symbol names provider absolutely needs to
 * work there. This functionality is covered by package scope tests.
 */
@OptIn(ExperimentalTopLevelDeclarationsGenerationApi::class)
internal object DeclarationsInPackageProvider {
    internal fun getTopLevelClassifierNamesInPackageProvider(packageFqName: FqName, analysisSession: KaFirSession): Set<Name> {
        return buildSet {
            addAll(
                analysisSession.symbolNamesProvider.getTopLevelClassifierNamesInPackage(packageFqName)
                    ?: analysisSession.useSiteScopeDeclarationProvider.getTopLevelKotlinClassLikeDeclarationNamesInPackage(packageFqName)
            )

            when {
                analysisSession.targetPlatform.isJvm() -> {
                    val psiPackage = PsiPackageImpl(PsiManager.getInstance(analysisSession.project), packageFqName.asString())
                    forEachNonKotlinPsiElementFinder(analysisSession.project) { finder ->
                        finder.getClassNames(psiPackage, analysisSession.analysisScope)
                            .mapNotNullTo(this, Name::identifier)
                    }
                }
            }

            addAll(collectGeneratedTopLevelClassifiers(packageFqName, analysisSession.firSession))
        }
    }

    internal fun getTopLevelCallableNamesInPackageProvider(packageFqName: FqName, analysisSession: KaFirSession): Set<Name> {
        return buildSet {
            addAll(
                analysisSession.symbolNamesProvider.getTopLevelCallableNamesInPackage(packageFqName)
                    ?: analysisSession.useSiteScopeDeclarationProvider.getTopLevelCallableNamesInPackage(packageFqName)
            )
            addAll(collectGeneratedTopLevelCallables(packageFqName, analysisSession.firSession))
        }
    }

    private val KaFirSession.symbolNamesProvider: FirSymbolNamesProvider
        get() = firSession.symbolProvider.symbolNamesProvider

    private fun collectGeneratedTopLevelClassifiers(packageFqName: FqName, session: FirSession): Set<Name> {
        val declarationGenerators = session.extensionService.declarationGenerators

        val generatedTopLevelClassifiers = declarationGenerators
            .asSequence()
            .flatMap {
                // FIXME this function should be called only once during plugin's lifetime, so this usage is not really correct (2)
                it.getTopLevelClassIds()
            }
            .filter { it.packageFqName == packageFqName }
            .map { it.shortClassName }

        return generatedTopLevelClassifiers.toSet()
    }

    private fun collectGeneratedTopLevelCallables(packageFqName: FqName, session: FirSession): Set<Name> {
        val generators = session.extensionService.declarationGenerators

        val generatedTopLevelDeclarations = generators
            .asSequence()
            .flatMap {
                // FIXME this function should be called only once during plugin's lifetime, so this usage is not really correct (1)
                it.getTopLevelCallableIds()
            }
            .filter { it.packageName == packageFqName }
            .map { it.callableName }

        return generatedTopLevelDeclarations.toSet()
    }
}
