/*
 * 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.plugin.sandbox.fir.generators

import org.jetbrains.kotlin.GeneratedDeclarationKey
import org.jetbrains.kotlin.descriptors.ClassKind
import org.jetbrains.kotlin.fir.FirSession
import org.jetbrains.kotlin.fir.extensions.*
import org.jetbrains.kotlin.fir.extensions.predicate.LookupPredicate
import org.jetbrains.kotlin.fir.plugin.createNestedClass
import org.jetbrains.kotlin.fir.symbols.impl.FirClassLikeSymbol
import org.jetbrains.kotlin.fir.symbols.impl.FirClassSymbol
import org.jetbrains.kotlin.name.Name
import org.jetbrains.kotlin.plugin.sandbox.fir.fqn

/**
 * ```
 * fun foo() {
 *   @AddNestedGeneratedClass
 *   interface Some {
 *       class Generated // will be generated by this extension
 *   }
 * }
 * ```
 */
class NestedClassGeneratorWithLocalClassesSupport(session: FirSession) : FirDeclarationGenerationExtension(session) {
    companion object {
        private val GENERATED = Name.identifier("Generated")
        private val ANNOTATION = "AddNestedGeneratedClass".fqn()
        private val PREDICATE = LookupPredicate.create {
            annotated(ANNOTATION)
        }
    }

    override fun getNestedClassifiersNames(classSymbol: FirClassSymbol<*>, context: NestedClassGenerationContext): Set<Name> {
        return when {
            session.predicateBasedProvider.matches(PREDICATE, classSymbol) -> setOf(GENERATED)
            else -> emptySet()
        }
    }

    override fun generateNestedClassLikeDeclaration(
        owner: FirClassSymbol<*>,
        name: Name,
        context: NestedClassGenerationContext,
    ): FirClassLikeSymbol<*>? {
        return when (name) {
            GENERATED -> {
                if (!session.predicateBasedProvider.matches(PREDICATE, owner)) return null
                createNestedClass(owner, GENERATED, Key, classKind = ClassKind.CLASS).symbol
            }
            else -> null
        }
    }

    override fun FirDeclarationPredicateRegistrar.registerPredicates() {
        register(PREDICATE)
    }

    private object Key : GeneratedDeclarationKey()
}
