/*
 * 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.low.level.api.fir.util

import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.llFirResolvableSession
import org.jetbrains.kotlin.fir.FirElementWithResolveState
import org.jetbrains.kotlin.fir.declarations.*
import org.jetbrains.kotlin.fir.psi
import org.jetbrains.kotlin.fir.resolve.providers.firProvider
import org.jetbrains.kotlin.psi.KtFile

fun FirElementWithResolveState.getContainingFile(): FirFile? {
    val provider = moduleData.session.firProvider
    return when (this) {
        is FirFile -> this
        is FirScript -> provider.getFirScriptContainerFile(symbol)
        is FirTypeParameter -> containingDeclarationSymbol.fir.getContainingFile()
        is FirPropertyAccessor -> propertySymbol.fir.getContainingFile()
        is FirValueParameter -> containingDeclarationSymbol.fir.getContainingFile()
        is FirBackingField -> propertySymbol.fir.getContainingFile()
        is FirCallableDeclaration -> provider.getFirCallableContainerFile(symbol)
        is FirClassLikeDeclaration -> provider.getFirClassifierContainerFileIfAny(symbol)
        is FirAnonymousInitializer -> {
            val classId = containingClassIdOrNull()
            if (classId?.isLocal == true) {
                containingKtFileIfAny?.let {
                    val moduleComponents = llFirResolvableSession?.moduleComponents
                    moduleComponents?.cache?.getCachedFirFile(it)
                }
            } else {
                containingDeclarationSymbol.fir.getContainingFile()
            }
        }
        is FirDanglingModifierList, is FirCodeFragment -> {
            val ktFile = psi?.containingFile as? KtFile
                ?: error("File for dangling modifier list cannot be null")
            val moduleComponents = llFirResolvableSession?.moduleComponents
                ?: error("LLFirResolvableModuleSession for dangling modifier list cannot be null")
            moduleComponents.cache.getCachedFirFile(ktFile)
                ?: error("Fir file for dandling modifier list cannot be null")
        }
        is FirReceiverParameter -> containingDeclarationSymbol.fir.getContainingFile()
        else -> errorWithFirSpecificEntries("Unsupported declaration ${this::class}", fir = this)
    }
}
