// Copyright 2000-2022 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file.

package org.jetbrains.kotlin.idea.codeInliner

import com.intellij.openapi.application.ModalityState
import com.intellij.openapi.diagnostic.ControlFlowException
import com.intellij.openapi.diagnostic.Logger
import com.intellij.openapi.progress.ProgressIndicator
import com.intellij.openapi.progress.ProgressManager
import com.intellij.openapi.progress.Task
import com.intellij.openapi.util.Key
import com.intellij.openapi.util.NlsContexts
import com.intellij.psi.PsiElement
import com.intellij.psi.search.GlobalSearchScope
import com.intellij.psi.search.searches.ReferencesSearch
import com.intellij.util.ModalityUiUtil
import org.jetbrains.kotlin.descriptors.DeclarationDescriptor
import org.jetbrains.kotlin.descriptors.FunctionDescriptor
import org.jetbrains.kotlin.idea.caches.resolve.resolveToCall
import org.jetbrains.kotlin.idea.caches.resolve.resolveToDescriptorIfAny
import org.jetbrains.kotlin.idea.core.compareDescriptors
import org.jetbrains.kotlin.idea.core.targetDescriptors
import org.jetbrains.kotlin.idea.intentions.ConvertReferenceToLambdaIntention
import org.jetbrains.kotlin.idea.intentions.SpecifyExplicitLambdaSignatureIntention
import org.jetbrains.kotlin.idea.references.KtSimpleReference
import org.jetbrains.kotlin.idea.base.projectStructure.scope.KotlinSourceFilterScope
import org.jetbrains.kotlin.idea.base.util.reformatted
import org.jetbrains.kotlin.idea.util.application.executeWriteCommand
import com.intellij.openapi.application.runReadAction
import org.jetbrains.kotlin.psi.*
import org.jetbrains.kotlin.psi.psiUtil.*
import org.jetbrains.kotlin.utils.addToStdlib.safeAs

interface UsageReplacementStrategy {
    fun createReplacer(usage: KtReferenceExpression): (() -> KtElement?)?

    companion object {
        val KEY = Key<Unit>("UsageReplacementStrategy.replaceUsages")
    }
}

private val LOG = Logger.getInstance(UsageReplacementStrategy::class.java)

fun UsageReplacementStrategy.replaceUsagesInWholeProject(
    targetPsiElement: PsiElement,
    @NlsContexts.DialogTitle progressTitle: String,
    @NlsContexts.Command commandName: String,
    unwrapSpecialUsages: Boolean = true,
) {
    val project = targetPsiElement.project
    ProgressManager.getInstance().run(
        object : Task.Modal(project, progressTitle, true) {
            override fun run(indicator: ProgressIndicator) {
                val usages = runReadAction {
                    val searchScope = KotlinSourceFilterScope.projectSources(GlobalSearchScope.projectScope(project), project)
                    ReferencesSearch.search(targetPsiElement, searchScope)
                        .filterIsInstance<KtSimpleReference<KtReferenceExpression>>()
                        .map { ref -> ref.expression }
                }

              ModalityUiUtil.invokeLaterIfNeeded(ModalityState.NON_MODAL) {
                  project.executeWriteCommand(commandName) {
                      this@replaceUsagesInWholeProject.replaceUsages(usages, unwrapSpecialUsages)
                  }
              }
            }
        })
}

fun UsageReplacementStrategy.replaceUsages(usages: Collection<KtReferenceExpression>, unwrapSpecialUsages: Boolean = true) {
    val usagesByFile = usages.groupBy { it.containingFile }

    for ((file, usagesInFile) in usagesByFile) {
        usagesInFile.forEach { it.putCopyableUserData(UsageReplacementStrategy.KEY, Unit) }

        // we should delete imports later to not affect other usages
        val importsToDelete = mutableListOf<KtImportDirective>()

        var usagesToProcess = usagesInFile
        while (usagesToProcess.isNotEmpty()) {
            if (processUsages(usagesToProcess, importsToDelete, unwrapSpecialUsages)) break

            // some usages may get invalidated we need to find them in the tree
            usagesToProcess = file.collectDescendantsOfType { it.getCopyableUserData(UsageReplacementStrategy.KEY) != null }
        }

        file.forEachDescendantOfType<KtSimpleNameExpression> { it.putCopyableUserData(UsageReplacementStrategy.KEY, null) }

        importsToDelete.forEach { it.delete() }
    }
}

/**
 * @return false if some usages were invalidated
 */
private fun UsageReplacementStrategy.processUsages(
    usages: List<KtReferenceExpression>,
    importsToDelete: MutableList<KtImportDirective>,
    unwrapSpecialUsages: Boolean,
): Boolean {
    val sortedUsages = usages.sortedWith { element1, element2 ->
        if (element1.parent.textRange.intersects(element2.parent.textRange)) {
            compareValuesBy(element2, element1) { it.startOffset }
        } else {
            compareValuesBy(element1, element2) { it.startOffset }
        }
    }

    var invalidUsagesFound = false
    for (usage in sortedUsages) {
        try {
            if (!usage.isValid) {
                invalidUsagesFound = true
                continue
            }

            if (unwrapSpecialUsages) {
                val specialUsage = unwrapSpecialUsageOrNull(usage)
                if (specialUsage != null) {
                    createReplacer(specialUsage)?.invoke()
                    continue
                }
            }

            //TODO: keep the import if we don't know how to replace some of the usages
            val importDirective = usage.getStrictParentOfType<KtImportDirective>()
            if (importDirective != null) {
                if (!importDirective.isAllUnder && importDirective.targetDescriptors().size == 1) {
                    importsToDelete.add(importDirective)
                }
                continue
            }

            createReplacer(usage)?.invoke()?.parent?.parent?.parent?.reformatted(true)
        } catch (e: Throwable) {
            if (e is ControlFlowException) throw e
            LOG.error(e)
        }
    }
    return !invalidUsagesFound
}

fun unwrapSpecialUsageOrNull(
    usage: KtReferenceExpression
): KtSimpleNameExpression? {
    if (usage !is KtSimpleNameExpression) return null

    when (val usageParent = usage.parent) {
        is KtCallableReferenceExpression -> {
            if (usageParent.callableReference != usage) return null
            val (name, descriptor) = usage.nameAndDescriptor
            return ConvertReferenceToLambdaIntention.applyTo(usageParent)?.let {
                findNewUsage(it, name, descriptor)
            }
        }

        is KtCallElement -> {
            for (valueArgument in usageParent.valueArguments.asReversed()) {
                val callableReferenceExpression = valueArgument.getArgumentExpression() as? KtCallableReferenceExpression ?: continue
                ConvertReferenceToLambdaIntention.applyTo(callableReferenceExpression)
            }

            val lambdaExpressions = usageParent.valueArguments.mapNotNull { it.getArgumentExpression() as? KtLambdaExpression }
            if (lambdaExpressions.isEmpty()) return null

            val (name, descriptor) = usage.nameAndDescriptor
            val grandParent = usageParent.parent
            for (lambdaExpression in lambdaExpressions) {
                val functionDescriptor = lambdaExpression.functionLiteral.resolveToDescriptorIfAny() as? FunctionDescriptor ?: continue
                if (functionDescriptor.valueParameters.isNotEmpty()) {
                    SpecifyExplicitLambdaSignatureIntention.applyTo(lambdaExpression)
                }
            }

            return grandParent.safeAs<KtElement>()?.let {
                findNewUsage(it, name, descriptor)
            }
        }

    }

    return null
}

private val KtSimpleNameExpression.nameAndDescriptor get() = getReferencedName() to resolveToCall()?.candidateDescriptor

private fun findNewUsage(
    element: KtElement,
    targetName: String?,
    targetDescriptor: DeclarationDescriptor?
): KtSimpleNameExpression? = element.findDescendantOfType {
    it.getReferencedName() == targetName && compareDescriptors(it.project, targetDescriptor, it.resolveToCall()?.candidateDescriptor)
}
