Get gradle version and gradle home from corresponding BuildModel after import

GradleProjectSettings that were used, are updated after KotlinDslListener is called, so it isn't correct to check if kotlinDslModels are supported using it. Also GradleScriptDefinitionsContributor should use gradle home from BuildModel, not from settings for the same reasons. ^KT-39104 Fixed
2020-05-29 17:37:28 +03:00
parent ece61915de
commit 7a47994498
19 changed files with 572 additions and 320 deletions
@@ -10,14 +10,19 @@ package org.jetbrains.kotlin.idea.scripting.gradle
 import com.intellij.notification.*
 import com.intellij.openapi.actionSystem.AnAction
 import com.intellij.openapi.actionSystem.AnActionEvent
+import com.intellij.openapi.application.ApplicationManager
 import com.intellij.openapi.externalSystem.importing.ImportSpecBuilder
 import com.intellij.openapi.externalSystem.util.ExternalSystemApiUtil
 import com.intellij.openapi.externalSystem.util.ExternalSystemUtil
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.util.Key
 import com.intellij.openapi.util.registry.Registry
+import com.intellij.openapi.vfs.VirtualFile
+import org.gradle.util.GradleVersion
 import org.jetbrains.kotlin.idea.KotlinIcons
 import org.jetbrains.kotlin.idea.KotlinIdeaGradleBundle
+import org.jetbrains.kotlin.idea.scripting.gradle.roots.GradleBuildRootsManager
+import org.jetbrains.kotlin.idea.util.application.getServiceSafe
 import org.jetbrains.kotlin.psi.UserDataProperty
 import org.jetbrains.plugins.gradle.settings.GradleProjectSettings
 import org.jetbrains.plugins.gradle.util.GradleConstants
@@ -37,26 +42,24 @@ fun runPartialGradleImport(project: Project) {

 fun getMissingConfigurationActionText() = KotlinIdeaGradleBundle.message("action.label.import.project")

-fun autoReloadScriptConfigurations(project: Project): Boolean {
-    val gradleSettings = ExternalSystemApiUtil.getSettings(project, GradleConstants.SYSTEM_ID)
-    val projectSettings = gradleSettings.getLinkedProjectsSettings()
-        .filterIsInstance<GradleProjectSettings>()
-        .firstOrNull()
-    if (projectSettings != null) {
-        return projectSettings.isUseAutoImport
-    }
+fun autoReloadScriptConfigurations(project: Project, file: VirtualFile): Boolean {
+    val path = GradleBuildRootsManager.getInstance(project)
+        .getScriptInfo(file)?.buildRoot?.pathPrefix ?: return false

-    return false
+    return ExternalSystemApiUtil
+        .getSettings(project, GradleConstants.SYSTEM_ID)
+        .getLinkedProjectSettings(path)
+        ?.isUseAutoImport ?: false
 }

 private const val kotlinDslNotificationGroupId = "Gradle Kotlin DSL Scripts"
 private var Project.notificationPanel: ScriptConfigurationChangedNotification?
        by UserDataProperty<Project, ScriptConfigurationChangedNotification>(Key.create("load.script.configuration.panel"))

-fun scriptConfigurationsNeedToBeUpdated(project: Project) {
+fun scriptConfigurationsNeedToBeUpdated(project: Project, file: VirtualFile) {
    if (!scriptConfigurationsNeedToBeUpdatedBalloon) return

-    if (autoReloadScriptConfigurations(project)) {
+    if (autoReloadScriptConfigurations(project, file)) {
        // import should be run automatically by Gradle plugin
        return
    }
@@ -98,12 +101,8 @@ private class ScriptConfigurationChangedNotification(val project: Project) :
    init {
        addAction(LoadConfigurationAction())
        addAction(NotificationAction.createSimple(KotlinIdeaGradleBundle.message("action.label.enable.auto.import")) {
-            val gradleSettings = ExternalSystemApiUtil.getSettings(project, GradleConstants.SYSTEM_ID)
-            val projectSettings = gradleSettings.getLinkedProjectsSettings()
-                .filterIsInstance<GradleProjectSettings>()
-                .firstOrNull()
-            if (projectSettings != null) {
-                projectSettings.isUseAutoImport = true
+            getGradleProjectSettings(project).forEach {
+                it.isUseAutoImport = true
            }
            runPartialGradleImport(project)
        })
@@ -25,6 +25,7 @@ import org.jetbrains.kotlin.idea.core.script.settings.KotlinScriptingSettings
 import org.jetbrains.kotlin.idea.scripting.gradle.importing.KotlinDslScriptModelResolver
 import org.jetbrains.kotlin.idea.scripting.gradle.roots.GradleBuildRootsManager
 import org.jetbrains.plugins.gradle.service.project.GradlePartialResolverPolicy
+import org.jetbrains.plugins.gradle.settings.GradleProjectSettings
 import org.jetbrains.plugins.gradle.util.GradleConstants

 fun runPartialGradleImport(project: Project) {
@@ -39,17 +40,17 @@ fun runPartialGradleImport(project: Project) {
    }
 }

-fun getMissingConfigurationNotificationText() = KotlinIdeaGradleBundle.message("script.configurations.will.be.available.after.load.changes")
 fun getMissingConfigurationActionText() = KotlinIdeaGradleBundle.message("action.text.load.script.configurations")

-fun autoReloadScriptConfigurations(project: Project): Boolean {
-    return GradleScriptDefinitionsContributor.getDefinitions(project).any {
+fun autoReloadScriptConfigurations(project: Project, file: VirtualFile): Boolean {
+    val buildRoot = GradleBuildRootsManager.getInstance(project).getScriptInfo(file)?.buildRoot ?: return false
+    return GradleScriptDefinitionsContributor.getDefinitions(project, buildRoot.pathPrefix, buildRoot.data.gradleHome)?.any {
        KotlinScriptingSettings.getInstance(project).autoReloadConfigurations(it)
-    }
+    } ?: return false
 }

-fun scriptConfigurationsNeedToBeUpdated(project: Project) {
-    if (autoReloadScriptConfigurations(project)) {
+fun scriptConfigurationsNeedToBeUpdated(project: Project, file: VirtualFile) {
+    if (autoReloadScriptConfigurations(project, file)) {
        runPartialGradleImport(project)
    } else {
        // notification is shown in LoadConfigurationAction
@@ -86,7 +87,7 @@ class LoadConfigurationAction : AnAction(
        val project = editor.project ?: return false
        val file = getKotlinScriptFile(editor) ?: return false

-        if (autoReloadScriptConfigurations(project)) {
+        if (autoReloadScriptConfigurations(project, file)) {
            return false
        }

@@ -94,13 +95,13 @@ class LoadConfigurationAction : AnAction(
    }

    private fun getKotlinScriptFile(editor: Editor): VirtualFile? {
-            return FileDocumentManager.getInstance()
-                .getFile(editor.document)
-                ?.takeIf {
-                    it !is LightVirtualFileBase
-                            && it.isValid
-                            && it.fileType != KotlinFileType.INSTANCE
-                            && isGradleKotlinScript(it)
-                }
-        }
+        return FileDocumentManager.getInstance()
+            .getFile(editor.document)
+            ?.takeIf {
+                it !is LightVirtualFileBase
+                        && it.isValid
+                        && it.fileType == KotlinFileType.INSTANCE
+                        && isGradleKotlinScript(it)
+            }
+    }
 }
@@ -6,29 +6,33 @@
 package org.jetbrains.kotlin.idea.scripting.gradle

 import com.intellij.execution.configurations.CommandLineTokenizer
-import com.intellij.openapi.diagnostic.Logger
 import com.intellij.openapi.externalSystem.util.ExternalSystemApiUtil
 import com.intellij.openapi.project.Project
 import com.intellij.util.EnvironmentUtil
+import com.intellij.util.PathUtil
+import org.gradle.util.GradleVersion
 import org.jetbrains.kotlin.idea.KotlinIdeaGradleBundle
 import org.jetbrains.kotlin.idea.core.script.ScriptDefinitionContributor
 import org.jetbrains.kotlin.idea.core.script.ScriptDefinitionSourceAsContributor
 import org.jetbrains.kotlin.idea.core.script.ScriptDefinitionsManager
 import org.jetbrains.kotlin.idea.core.script.loadDefinitionsFromTemplates
 import org.jetbrains.kotlin.idea.scripting.gradle.importing.KotlinDslSyncListener
+import org.jetbrains.kotlin.idea.scripting.gradle.roots.GradleBuildRootsManager
+import org.jetbrains.kotlin.idea.scripting.gradle.roots.Imported
+import org.jetbrains.kotlin.idea.scripting.gradle.roots.WithoutScriptModels
 import org.jetbrains.kotlin.scripting.definitions.KotlinScriptDefinitionAdapterFromNewAPIBase
 import org.jetbrains.kotlin.scripting.definitions.ScriptDefinition
 import org.jetbrains.kotlin.scripting.definitions.getEnvironment
 import org.jetbrains.kotlin.scripting.resolve.KotlinScriptDefinitionFromAnnotatedTemplate
 import org.jetbrains.plugins.gradle.config.GradleSettingsListenerAdapter
+import org.jetbrains.plugins.gradle.service.GradleInstallationManager
 import org.jetbrains.plugins.gradle.settings.DistributionType
 import org.jetbrains.plugins.gradle.settings.GradleExecutionSettings
-import org.jetbrains.plugins.gradle.settings.GradleProjectSettings
 import org.jetbrains.plugins.gradle.settings.GradleSettingsListener
 import org.jetbrains.plugins.gradle.util.GradleConstants
 import java.io.File
 import java.util.*
-import java.util.concurrent.atomic.AtomicBoolean
+import java.util.concurrent.ConcurrentHashMap
 import kotlin.reflect.KClass
 import kotlin.script.dependencies.Environment
 import kotlin.script.dependencies.ScriptContents
@@ -42,23 +46,236 @@ import kotlin.script.templates.standard.ScriptTemplateWithArgs

 class GradleScriptDefinitionsContributor(private val project: Project) : ScriptDefinitionSourceAsContributor {
    companion object {
-        fun getDefinitions(project: Project) =
-            ScriptDefinitionContributor.EP_NAME.getExtensions(project)
+        fun getDefinitions(project: Project, workingDir: String, gradleHome: String?): List<ScriptDefinition>? {
+            val contributor = ScriptDefinitionContributor.EP_NAME.getExtensions(project)
                .filterIsInstance<GradleScriptDefinitionsContributor>()
-                .single().definitions.toList()
+                .singleOrNull()

-        fun getDefinitionsTemplateClasspath(project: Project): MutableSet<String> =
-            mutableSetOf<String>().also { result ->
-                getDefinitions(project).forEach { definition ->
-                    definition.asLegacyOrNull<KotlinScriptDefinitionFromAnnotatedTemplate>()
-                        ?.templateClasspath
-                        ?.mapNotNullTo(result) { it.path }
-                }
+            if (contributor == null) {
+                scriptingInfoLog(
+                    "cannot find gradle script definitions contributor in ScriptDefinitionContributor.EP_NAME list: " +
+                            "workingDir=$workingDir gradleHome=$gradleHome"
+                )
+                return null
            }
+            if (gradleHome == null) {
+                scriptingInfoLog(KotlinIdeaGradleBundle.message("error.text.unable.to.get.gradle.home.directory") + ": workingDir=$workingDir gradleHome=$gradleHome")
+                return null
+            }
+
+            val root = LightGradleBuildRoot(workingDir, gradleHome)
+            val definitions = contributor.definitionsByRoots[root]
+            if (definitions == null) {
+                scriptingInfoLog(
+                    "script definitions aren't loaded yet. " +
+                            "They should be loaded by invoking GradleScriptDefinitionsContributor.reloadIfNeeded from KotlinDslSyncListener: " +
+                            "workingDir=$workingDir gradleHome=$gradleHome"
+                )
+                return null
+            }
+            return definitions
+        }
+
+        fun getDefinitionsTemplateClasspath(gradleHome: String?): List<String> {
+            return try {
+                getFullDefinitionsClasspath(gradleHome).first.map {
+                    PathUtil.toSystemIndependentName(it.path)
+                }
+            } catch (e: Throwable) {
+                scriptingInfoLog("cannot get gradle classpath for Gradle Kotlin DSL scripts: ${e.message}")
+
+                emptyList()
+            }
+        }
+
+        private val kotlinDslDependencySelector = Regex("^gradle-(?:kotlin-dsl|core).*\\.jar\$")
+
+        private fun getFullDefinitionsClasspath(gradleHome: String?): Pair<List<File>, List<File>> {
+            if (gradleHome == null) {
+                error(KotlinIdeaGradleBundle.message("error.text.unable.to.get.gradle.home.directory"))
+            }
+
+            val gradleLibDir = File(gradleHome, "lib")
+                .let {
+                    it.takeIf { it.exists() && it.isDirectory }
+                        ?: error(KotlinIdeaGradleBundle.message("error.text.invalid.gradle.libraries.directory", it))
+                }
+
+            val templateClasspath = gradleLibDir
+                /* an inference problem without explicit 'it', TODO: remove when fixed */
+                .listFiles { it -> kotlinDslDependencySelector.matches(it.name) }
+                ?.takeIf { it.isNotEmpty() }
+                ?.asList()
+                ?: error(KotlinIdeaGradleBundle.message("error.text.missing.jars.in.gradle.directory"))
+
+            scriptingDebugLog { "gradle script templates classpath $templateClasspath" }
+
+            val additionalClassPath = kotlinStdlibAndCompiler(gradleLibDir)
+
+            scriptingDebugLog { "gradle script templates additional classpath $templateClasspath" }
+
+            return templateClasspath to additionalClassPath
+        }
+
+        // TODO: check this against kotlin-dsl branch that uses daemon
+        private fun kotlinStdlibAndCompiler(gradleLibDir: File): List<File> {
+            // additionally need compiler jar to load gradle resolver
+            return gradleLibDir.listFiles { file ->
+                file?.name?.startsWith("kotlin-compiler-embeddable") == true || file?.name?.startsWith("kotlin-stdlib") == true
+            }?.firstOrNull()?.let(::listOf).orEmpty()
+        }
    }

    override val id: String = "Gradle Kotlin DSL"
-    private val failedToLoad = AtomicBoolean(false)
+
+    private data class LightGradleBuildRoot(val workingDir: String, val gradleHome: String?)
+
+    private val definitionsByRoots = ConcurrentHashMap<LightGradleBuildRoot, List<ScriptDefinition>>()
+
+    private fun LightGradleBuildRoot.markAsError() {
+        definitionsByRoots[this] = listOf(ErrorGradleScriptDefinition(project))
+    }
+
+    private fun LightGradleBuildRoot.isError(): Boolean {
+        return definitionsByRoots[this]?.any { it is ErrorGradleScriptDefinition } ?: false
+    }
+
+    private fun forceReload() {
+        for (key in definitionsByRoots.keys()) {
+            key.markAsError()
+        }
+        ScriptDefinitionsManager.getInstance(project).reloadDefinitionsBy(this)
+    }
+
+    // TODO: remove old roots
+    fun reloadIfNeeded(workingDir: String, gradleHome: String?) {
+        val root = LightGradleBuildRoot(workingDir, gradleHome)
+        val value = definitionsByRoots[root]
+        if (value != null) {
+            if (root.isError()) {
+                ScriptDefinitionsManager.getInstance(project).reloadDefinitionsBy(this)
+            }
+        } else {
+            val rootWithChangedGradleHome = definitionsByRoots.filter { it.key.workingDir == workingDir }
+            if (rootWithChangedGradleHome.isNotEmpty()) {
+                rootWithChangedGradleHome.forEach {
+                    definitionsByRoots.remove(it.key)
+                }
+            }
+            root.markAsError()
+            ScriptDefinitionsManager.getInstance(project).reloadDefinitionsBy(this)
+        }
+    }
+
+    private fun loadGradleDefinitions(root: LightGradleBuildRoot): List<ScriptDefinition> {
+        try {
+            val (templateClasspath, additionalClassPath) = getFullDefinitionsClasspath(root.gradleHome)
+
+            val kotlinDslTemplates = ArrayList<ScriptDefinition>()
+
+            val projectPath = root.workingDir
+            loadGradleTemplates(
+                projectPath,
+                templateClass = "org.gradle.kotlin.dsl.KotlinInitScript",
+                root.gradleHome,
+                templateClasspath,
+                additionalClassPath
+            ).let { kotlinDslTemplates.addAll(it) }
+
+            loadGradleTemplates(
+                projectPath,
+                templateClass = "org.gradle.kotlin.dsl.KotlinSettingsScript",
+                gradleHome = root.gradleHome,
+                templateClasspath = templateClasspath,
+                additionalClassPath = additionalClassPath
+            ).let { kotlinDslTemplates.addAll(it) }
+
+            // KotlinBuildScript should be last because it has wide scriptFilePattern
+            loadGradleTemplates(
+                projectPath,
+                templateClass = "org.gradle.kotlin.dsl.KotlinBuildScript",
+                gradleHome = root.gradleHome,
+                templateClasspath = templateClasspath,
+                additionalClassPath = additionalClassPath
+            ).let { kotlinDslTemplates.addAll(it) }
+
+
+            if (kotlinDslTemplates.isNotEmpty()) {
+                return kotlinDslTemplates.distinct()
+            }
+        } catch (t: Throwable) {
+            // TODO: review exception handling
+
+            if (t is IllegalStateException) {
+                scriptingInfoLog("IllegalStateException loading gradle script templates: ${t.message}")
+            } else {
+                scriptingDebugLog { "error loading gradle script templates ${t.message}" }
+            }
+
+            return listOf(ErrorGradleScriptDefinition(project, t.message))
+        }
+
+        return listOf(ErrorGradleScriptDefinition(project))
+
+    }
+
+    private fun loadGradleTemplates(
+        projectPath: String,
+        templateClass: String,
+        gradleHome: String?,
+        templateClasspath: List<File>,
+        additionalClassPath: List<File>
+    ): List<ScriptDefinition> {
+        val gradleExeSettings = ExternalSystemApiUtil.getExecutionSettings<GradleExecutionSettings>(
+            project,
+            projectPath,
+            GradleConstants.SYSTEM_ID
+        )
+        val hostConfiguration = createHostConfiguration(projectPath, gradleHome, gradleExeSettings)
+        return loadDefinitionsFromTemplates(
+            listOf(templateClass),
+            templateClasspath,
+            hostConfiguration,
+            additionalClassPath
+        ).map {
+            it.asLegacyOrNull<KotlinScriptDefinitionFromAnnotatedTemplate>()?.let { legacyDef ->
+                // Expand scope for old gradle script definition
+                val version = GradleInstallationManager.getGradleVersion(gradleHome) ?: GradleVersion.current().version
+                GradleKotlinScriptDefinitionWrapper(
+                    it.hostConfiguration,
+                    legacyDef,
+                    version
+                )
+            } ?: it
+        }
+    }
+
+    private fun createHostConfiguration(
+        projectPath: String,
+        gradleHome: String?,
+        gradleExeSettings: GradleExecutionSettings
+    ): ScriptingHostConfiguration {
+        val gradleJvmOptions = gradleExeSettings.daemonVmOptions?.let { vmOptions ->
+            CommandLineTokenizer(vmOptions).toList()
+                .mapNotNull { it?.let { it as? String } }
+                .filterNot(String::isBlank)
+                .distinct()
+        } ?: emptyList()
+
+        val environment = mapOf(
+            "gradleHome" to gradleHome?.let(::File),
+            "gradleJavaHome" to gradleExeSettings.javaHome, // TODO: should be taken from information from gradle sync
+
+            "projectRoot" to projectPath.let(::File),
+
+            "gradleOptions" to emptyList<String>(), // There is no option in UI to set project wide gradleOptions
+            "gradleJvmOptions" to gradleJvmOptions,
+            "gradleEnvironmentVariables" to if (gradleExeSettings.isPassParentEnvs) EnvironmentUtil.getEnvironmentMap() else emptyMap()
+        )
+        return ScriptingHostConfiguration(defaultJvmScriptingHostConfiguration) {
+            getEnvironment { environment }
+        }
+    }

    init {
        subscribeToGradleSettingChanges()
@@ -67,19 +284,19 @@ class GradleScriptDefinitionsContributor(private val project: Project) : ScriptD
    private fun subscribeToGradleSettingChanges() {
        val listener = object : GradleSettingsListenerAdapter() {
            override fun onGradleVmOptionsChange(oldOptions: String?, newOptions: String?) {
-                reload()
+                forceReload()
            }

            override fun onGradleHomeChange(oldPath: String?, newPath: String?, linkedProjectPath: String) {
-                reload()
+                forceReload()
            }

            override fun onServiceDirectoryPathChange(oldPath: String?, newPath: String?) {
-                reload()
+                forceReload()
            }

            override fun onGradleDistributionTypeChange(currentValue: DistributionType?, linkedProjectPath: String) {
-                reload()
+                forceReload()
            }
        }
        project.messageBus.connect().subscribe(GradleSettingsListener.TOPIC, listener)
@@ -89,214 +306,87 @@ class GradleScriptDefinitionsContributor(private val project: Project) : ScriptD
    // TODO: possibly combine exceptions from every loadGradleTemplates call, be mindful of KT-19276
    override val definitions: Sequence<ScriptDefinition>
        get() {
-            val kotlinDslDependencySelector = Regex("^gradle-(?:kotlin-dsl|core).*\\.jar\$")
-            val kotlinDslAdditionalResolverCp = ::kotlinStdlibAndCompiler
-
-            failedToLoad.set(false)
-
-            val kotlinDslTemplates = ArrayList<ScriptDefinition>()
-
-            loadGradleTemplates(
-                templateClass = "org.gradle.kotlin.dsl.KotlinInitScript",
-                dependencySelector = kotlinDslDependencySelector,
-                additionalResolverClasspath = kotlinDslAdditionalResolverCp
-
-            ).let { kotlinDslTemplates.addAll(it) }
-
-            loadGradleTemplates(
-                templateClass = "org.gradle.kotlin.dsl.KotlinSettingsScript",
-                dependencySelector = kotlinDslDependencySelector,
-                additionalResolverClasspath = kotlinDslAdditionalResolverCp
-
-            ).let { kotlinDslTemplates.addAll(it) }
-
-            // KotlinBuildScript should be last because it has wide scriptFilePattern
-            loadGradleTemplates(
-                templateClass = "org.gradle.kotlin.dsl.KotlinBuildScript",
-                dependencySelector = kotlinDslDependencySelector,
-                additionalResolverClasspath = kotlinDslAdditionalResolverCp
-            ).let { kotlinDslTemplates.addAll(it) }
-
-
-            if (kotlinDslTemplates.isNotEmpty()) {
-                return kotlinDslTemplates.distinct().asSequence()
+            definitionsByRoots.keys().iterator().forEachRemaining { root ->
+                // reload definitions marked as error
+                if (root.isError()) {
+                    definitionsByRoots[root] = loadGradleDefinitions(root)
+                }
            }
-
-            val default = tryToLoadOldBuildScriptDefinition()
-            if (default.isNotEmpty()) {
-                return default.distinct().asSequence()
+            if (definitionsByRoots.isEmpty()) {
+                // can be empty in case when import wasn't done from IDE start up,
+                // otherwise KotlinDslSyncListener should run reloadIfNeeded for valid roots
+                GradleBuildRootsManager.getInstance(project).getAllRoots().forEach {
+                    val workingDir = it.pathPrefix
+                    val gradleHome = when (it) {
+                        is Imported -> {
+                            it.data.gradleHome
+                        }
+                        is WithoutScriptModels -> {
+                            ExternalSystemApiUtil.getExecutionSettings<GradleExecutionSettings>(
+                                project,
+                                workingDir,
+                                GradleConstants.SYSTEM_ID
+                            ).gradleHome
+                        }
+                    }
+                    val root = LightGradleBuildRoot(workingDir, gradleHome)
+                    definitionsByRoots[root] = loadGradleDefinitions(root)
+                }
+            }
+            if (definitionsByRoots.isNotEmpty()) {
+                return definitionsByRoots.flatMap { it.value }.asSequence()
            }
-
            return sequenceOf(ErrorGradleScriptDefinition(project))
        }

-    private fun tryToLoadOldBuildScriptDefinition(): List<ScriptDefinition> {
-        failedToLoad.set(false)

-        return loadGradleTemplates(
-            templateClass = "org.gradle.script.lang.kotlin.KotlinBuildScript",
-            dependencySelector = Regex("^gradle-(?:script-kotlin|core).*\\.jar\$"),
-            additionalResolverClasspath = { emptyList() }
-        )
-    }
+}

-    // TODO: check this against kotlin-dsl branch that uses daemon
-    private fun kotlinStdlibAndCompiler(gradleLibDir: File): List<File> {
-        // additionally need compiler jar to load gradle resolver
-        return gradleLibDir.listFiles { file ->
-            file.name.startsWith("kotlin-compiler-embeddable") || file.name.startsWith("kotlin-stdlib")
+// TODO: refactor - minimize
+private class ErrorGradleScriptDefinition(project: Project, message: String? = null) :
+    ScriptDefinition.FromLegacy(
+        ScriptingHostConfiguration(defaultJvmScriptingHostConfiguration),
+        LegacyDefinition(project, message)
+    ) {
+
+    private class LegacyDefinition(project: Project, message: String?) : KotlinScriptDefinitionAdapterFromNewAPIBase() {
+        companion object {
+            private const val KOTLIN_DSL_SCRIPT_EXTENSION = "gradle.kts"
        }
-            .firstOrNull()?.let(::listOf).orEmpty()
+
+        override val name: String get() = KotlinIdeaGradleBundle.message("text.default.kotlin.gradle.script")
+        override val fileExtension: String = KOTLIN_DSL_SCRIPT_EXTENSION
+
+        override val scriptCompilationConfiguration: ScriptCompilationConfiguration = ScriptCompilationConfiguration.Default
+        override val hostConfiguration: ScriptingHostConfiguration = ScriptingHostConfiguration(defaultJvmScriptingHostConfiguration)
+        override val baseClass: KClass<*> = ScriptTemplateWithArgs::class
+
+        override val dependencyResolver: DependenciesResolver =
+            ErrorScriptDependenciesResolver(project, message)
+
+        override fun toString(): String = "ErrorGradleScriptDefinition"
    }

-    private fun loadGradleTemplates(
-        templateClass: String, dependencySelector: Regex,
-        additionalResolverClasspath: (gradleLibDir: File) -> List<File>
-    ): List<ScriptDefinition> = try {
-        doLoadGradleTemplates(templateClass, dependencySelector, additionalResolverClasspath)
-    } catch (t: Throwable) {
-        scriptingDebugLog { "error loading gradle script templates ${t.message}" }
-        // TODO: review exception handling
-        failedToLoad.set(true)
-        if (t is IllegalStateException) {
-            Logger.getInstance(GradleScriptDefinitionsContributor::class.java)
-                .info("[kts] error loading gradle script templates: ${t.message}")
-        }
-        listOf(ErrorGradleScriptDefinition(project, t.message))
-    }
+    override fun equals(other: Any?): Boolean = other is ErrorGradleScriptDefinition
+    override fun hashCode(): Int = name.hashCode()
+}

-
-    private fun doLoadGradleTemplates(
-        templateClass: String, dependencySelector: Regex,
-        additionalResolverClasspath: (gradleLibDir: File) -> List<File>
-    ): List<ScriptDefinition> {
-        fun createHostConfiguration(
-            gradleExeSettings: GradleExecutionSettings,
-            projectSettings: GradleProjectSettings
-        ): ScriptingHostConfiguration {
-            val gradleJvmOptions = gradleExeSettings.daemonVmOptions?.let { vmOptions ->
-                CommandLineTokenizer(vmOptions).toList()
-                    .mapNotNull { it?.let { it as? String } }
-                    .filterNot(String::isBlank)
-                    .distinct()
-            } ?: emptyList()
-
-
-            val environment = mapOf(
-                "gradleHome" to gradleExeSettings.gradleHome?.let(::File),
-                "gradleJavaHome" to gradleExeSettings.javaHome,
-
-                "projectRoot" to projectSettings.externalProjectPath.let(::File),
-
-                "gradleOptions" to emptyList<String>(), // There is no option in UI to set project wide gradleOptions
-                "gradleJvmOptions" to gradleJvmOptions,
-                "gradleEnvironmentVariables" to if (gradleExeSettings.isPassParentEnvs) EnvironmentUtil.getEnvironmentMap() else emptyMap()
+private class ErrorScriptDependenciesResolver(
+    private val project: Project,
+    private val message: String? = null
+) : DependenciesResolver {
+    override fun resolve(scriptContents: ScriptContents, environment: Environment): ResolveResult {
+        val importTasks = KotlinDslSyncListener.instance.tasks
+        val importInProgress = synchronized(importTasks) { importTasks.values.any { it.project == project } }
+        val failureMessage = if (importInProgress) {
+            KotlinIdeaGradleBundle.message("error.text.highlighting.is.impossible.during.gradle.import")
+        } else {
+            message ?: KotlinIdeaGradleBundle.message(
+                "error.text.failed.to.load.script.definitions.by",
+                GradleScriptDefinitionsContributor::class.java.name
            )
-            return ScriptingHostConfiguration(defaultJvmScriptingHostConfiguration) {
-                getEnvironment { environment }
-            }
-        }
-
-        val gradleSettings = ExternalSystemApiUtil.getSettings(project, GradleConstants.SYSTEM_ID)
-        if (gradleSettings.getLinkedProjectsSettings().isEmpty())
-            error(KotlinIdeaGradleBundle.message("error.text.project.isn.t.linked.with.gradle", project.name))
-
-        val projectSettings = gradleSettings.getLinkedProjectsSettings().filterIsInstance<GradleProjectSettings>().firstOrNull()
-            ?: error(KotlinIdeaGradleBundle.message("error.text.project.isn.t.linked.with.gradle", project.name))
-
-        val gradleExeSettings = ExternalSystemApiUtil.getExecutionSettings<GradleExecutionSettings>(
-            project,
-            projectSettings.externalProjectPath,
-            GradleConstants.SYSTEM_ID
-        )
-
-        val gradleHome = gradleExeSettings.gradleHome
-            ?: error(KotlinIdeaGradleBundle.message("error.text.unable.to.get.gradle.home.directory"))
-
-        val gradleLibDir = File(gradleHome, "lib").let {
-            it.takeIf { it.exists() && it.isDirectory }
-                ?: error(KotlinIdeaGradleBundle.message("error.text.invalid.gradle.libraries.directory", it))
-        }
-        val templateClasspath = gradleLibDir.listFiles { it ->
-            /* an inference problem without explicit 'it', TODO: remove when fixed */
-            dependencySelector.matches(it.name)
-        }.takeIf { it.isNotEmpty() }?.asList() ?: error(KotlinIdeaGradleBundle.message("error.text.missing.jars.in.gradle.directory"))
-
-        scriptingDebugLog { "loading gradle script templates from $templateClasspath" }
-
-        return loadDefinitionsFromTemplates(
-            listOf(templateClass),
-            templateClasspath,
-            createHostConfiguration(gradleExeSettings, projectSettings),
-            additionalResolverClasspath(gradleLibDir)
-        ).map {
-            it.asLegacyOrNull<KotlinScriptDefinitionFromAnnotatedTemplate>()?.let { legacyDef ->
-                // Expand scope for old gradle script definition
-                GradleKotlinScriptDefinitionWrapper(
-                    it.hostConfiguration,
-                    legacyDef,
-                    projectSettings.resolveGradleVersion().version
-                )
-            } ?: it
-        }
-    }
-
-    fun reloadIfNecessary() {
-        if (failedToLoad.compareAndSet(true, false)) {
-            reload()
-        }
-    }
-
-    private fun reload() {
-        ScriptDefinitionsManager.getInstance(project).reloadDefinitionsBy(this)
-    }
-
-    // TODO: refactor - minimize
-    private class ErrorGradleScriptDefinition(project: Project, message: String? = null) :
-        ScriptDefinition.FromLegacy(
-            ScriptingHostConfiguration(defaultJvmScriptingHostConfiguration),
-            LegacyDefinition(project, message)
-        ) {
-
-        private class LegacyDefinition(project: Project, message: String?) : KotlinScriptDefinitionAdapterFromNewAPIBase() {
-            companion object {
-                private const val KOTLIN_DSL_SCRIPT_EXTENSION = "gradle.kts"
-            }
-
-            override val name: String get() = KotlinIdeaGradleBundle.message("text.default.kotlin.gradle.script")
-            override val fileExtension: String = KOTLIN_DSL_SCRIPT_EXTENSION
-
-            override val scriptCompilationConfiguration: ScriptCompilationConfiguration = ScriptCompilationConfiguration.Default
-            override val hostConfiguration: ScriptingHostConfiguration = ScriptingHostConfiguration(defaultJvmScriptingHostConfiguration)
-            override val baseClass: KClass<*> = ScriptTemplateWithArgs::class
-
-            override val dependencyResolver: DependenciesResolver =
-                ErrorScriptDependenciesResolver(project, message)
-
-            override fun toString(): String = "ErrorGradleScriptDefinition"
-        }
-
-        override fun equals(other: Any?): Boolean = other is ErrorGradleScriptDefinition
-        override fun hashCode(): Int = name.hashCode()
-    }
-
-    private class ErrorScriptDependenciesResolver(
-        private val project: Project,
-        private val message: String? = null
-    ) : DependenciesResolver {
-        override fun resolve(scriptContents: ScriptContents, environment: Environment): ResolveResult {
-            val importTasks = KotlinDslSyncListener.instance.tasks
-            val importInProgress = synchronized(importTasks) { importTasks.values.any { it.project == project } }
-            val failureMessage = if (importInProgress) {
-                KotlinIdeaGradleBundle.message("error.text.highlighting.is.impossible.during.gradle.import")
-            } else {
-                message ?: KotlinIdeaGradleBundle.message(
-                    "error.text.failed.to.load.script.definitions.by",
-                    GradleScriptDefinitionsContributor::class.java.name
-                )
-            }
-            return ResolveResult.Failure(ScriptReport(failureMessage, ScriptReport.Severity.FATAL))
        }
+        return ResolveResult.Failure(ScriptReport(failureMessage, ScriptReport.Severity.FATAL))
    }
 }

@@ -7,7 +7,6 @@ package org.jetbrains.kotlin.idea.scripting.gradle.importing

 import com.intellij.openapi.externalSystem.model.task.ExternalSystemTaskId
 import com.intellij.openapi.project.Project
-import com.intellij.openapi.util.Key
 import com.intellij.openapi.util.io.FileUtil.toSystemIndependentName
 import com.intellij.openapi.vfs.VfsUtil
 import org.gradle.tooling.model.kotlin.dsl.EditorReportSeverity
@@ -16,10 +15,26 @@ import org.jetbrains.kotlin.gradle.BrokenKotlinDslScriptsModel
 import org.jetbrains.kotlin.idea.KotlinIdeaGradleBundle
 import org.jetbrains.kotlin.idea.scripting.gradle.getGradleScriptInputsStamp
 import org.jetbrains.kotlin.idea.scripting.gradle.roots.GradleBuildRootsManager
-import org.jetbrains.kotlin.psi.NotNullableUserDataProperty
+import org.jetbrains.plugins.gradle.model.BuildScriptClasspathModel
 import org.jetbrains.plugins.gradle.service.project.ProjectResolverContext
 import java.io.File
-import java.util.concurrent.ConcurrentHashMap
+
+fun saveGradleBuildEnvironment(resolverCtx: ProjectResolverContext) {
+    val task = resolverCtx.externalSystemTaskId
+    val tasks = KotlinDslSyncListener.instance.tasks
+    val sync = synchronized(tasks) { tasks[task] }
+    if (sync != null) {
+        val gradleHome = resolverCtx.getExtraProject(BuildScriptClasspathModel::class.java)?.gradleHomeDir?.canonicalPath
+            ?: resolverCtx.settings?.gradleHome
+        synchronized(sync) {
+            sync.gradleVersion = resolverCtx.projectGradleVersion
+
+            if (gradleHome != null) {
+                sync.gradleHome = toSystemIndependentName(gradleHome)
+            }
+        }
+    }
+}

 fun processScriptModel(
    resolverCtx: ProjectResolverContext,
@@ -44,6 +59,11 @@ fun processScriptModel(
        }

        if (models.containsErrors()) {
+            if (sync != null) {
+                synchronized(sync) {
+                    sync.failed = true
+                }
+            }
            throw IllegalStateException(KotlinIdeaGradleBundle.message("title.kotlin.build.script"))
        }
    }
@@ -103,6 +123,8 @@ private fun KotlinDslScriptsModel.toListOfScriptModels(project: Project): List<K
 class KotlinDslGradleBuildSync(val workingDir: String, val taskId: ExternalSystemTaskId) {
    val ts = System.currentTimeMillis()
    var project: Project? = null
+    var gradleVersion: String? = null
+    var gradleHome: String? = null
    val projectRoots = mutableSetOf<String>()
    val models = mutableListOf<KotlinDslScriptModel>()
    var failed = false
@@ -31,12 +31,10 @@ class KotlinDslScriptModelResolver : KotlinDslScriptModelResolverCommon() {
    override fun populateProjectExtraModels(gradleProject: IdeaProject, ideProject: DataNode<ProjectData>) {
        super.populateProjectExtraModels(gradleProject, ideProject)

-        if (kotlinDslScriptsModelImportSupported(resolverCtx.projectGradleVersion)) {
-            populateBuildModels(resolverCtx.models.mainBuild, ideProject)
+        populateBuildModels(resolverCtx.models.mainBuild, ideProject)

-            resolverCtx.models.includedBuilds.forEach { includedRoot ->
-                populateBuildModels(includedRoot, ideProject)
-            }
+        resolverCtx.models.includedBuilds.forEach { includedRoot ->
+            populateBuildModels(includedRoot, ideProject)
        }
    }

@@ -46,9 +44,13 @@ class KotlinDslScriptModelResolver : KotlinDslScriptModelResolverCommon() {
    ) {
        root.projects.forEach {
            if (it.projectIdentifier.projectPath == ":") {
-                resolverCtx.models.getModel(it, KotlinDslScriptsModel::class.java)?.let { model ->
-                    processScriptModel(resolverCtx, model, it.projectIdentifier.projectPath)
+                if (kotlinDslScriptsModelImportSupported(resolverCtx.projectGradleVersion)) {
+                    resolverCtx.models.getModel(it, KotlinDslScriptsModel::class.java)?.let { model ->
+                        processScriptModel(resolverCtx, model, it.projectIdentifier.projectPath)
+                    }
                }
+
+                saveGradleBuildEnvironment(resolverCtx)
            }
        }
    }
@@ -29,12 +29,10 @@ class KotlinDslScriptModelResolver : KotlinDslScriptModelResolverCommon() {
    override fun populateProjectExtraModels(gradleProject: IdeaProject, ideProject: DataNode<ProjectData>) {
        super.populateProjectExtraModels(gradleProject, ideProject)

-        if (kotlinDslScriptsModelImportSupported(resolverCtx.projectGradleVersion)) {
-            populateBuildModels(gradleProject, ideProject)
+        populateBuildModels(gradleProject, ideProject)

-            resolverCtx.models.includedBuilds.forEach { includedRoot ->
-                populateBuildModels(includedRoot, ideProject)
-            }
+        resolverCtx.models.includedBuilds.forEach { includedRoot ->
+            populateBuildModels(includedRoot, ideProject)
        }
    }

@@ -44,9 +42,13 @@ class KotlinDslScriptModelResolver : KotlinDslScriptModelResolverCommon() {
    ) {
        root.modules.forEach {
            if (it.gradleProject.parent == null) {
-                resolverCtx.getExtraProject(it, KotlinDslScriptsModel::class.java)?.let { model ->
-                    processScriptModel(resolverCtx, model, it.name)
+                if (kotlinDslScriptsModelImportSupported(resolverCtx.projectGradleVersion)) {
+                    resolverCtx.getExtraProject(it, KotlinDslScriptsModel::class.java)?.let { model ->
+                        processScriptModel(resolverCtx, model, it.name)
+                    }
                }
+
+                saveGradleBuildEnvironment(resolverCtx)
            }
        }
    }
@@ -5,12 +5,12 @@

 package org.jetbrains.kotlin.idea.scripting.gradle.importing

-import org.jetbrains.plugins.gradle.model.ClassSetImportModelProvider
-import org.jetbrains.plugins.gradle.model.ProjectImportModelProvider
+import org.gradle.tooling.model.kotlin.dsl.KotlinDslScriptsModel
 import org.jetbrains.kotlin.gradle.KotlinDslScriptAdditionalTask
 import org.jetbrains.kotlin.gradle.KotlinDslScriptModelProvider
-import org.gradle.tooling.model.kotlin.dsl.KotlinDslScriptsModel
 import org.jetbrains.kotlin.idea.scripting.gradle.kotlinDslScriptsModelImportSupported
+import org.jetbrains.plugins.gradle.model.ClassSetImportModelProvider
+import org.jetbrains.plugins.gradle.model.ProjectImportModelProvider
 import org.jetbrains.plugins.gradle.service.project.ModifiableGradleProjectModel
 import org.jetbrains.plugins.gradle.service.project.ProjectModelContributor
 import org.jetbrains.plugins.gradle.service.project.ProjectResolverContext
@@ -34,17 +34,19 @@ class KotlinDslScriptModelContributor : ProjectModelContributor {
    override fun accept(
        projectModelBuilder: ModifiableGradleProjectModel,
        toolingModelsProvider: ToolingModelsProvider,
-        resolverContext: ProjectResolverContext
+        resolverCtx: ProjectResolverContext
    ) {
-        if (!kotlinDslScriptsModelImportSupported(resolverContext.projectGradleVersion)) return
-
        toolingModelsProvider.projects().forEach {
            val projectIdentifier = it.projectIdentifier.projectPath
            if (projectIdentifier == ":") {
-                val model = toolingModelsProvider.getProjectModel(it, KotlinDslScriptsModel::class.java)
-                if (model != null) {
-                    processScriptModel(resolverContext, model, projectIdentifier)
+                if (kotlinDslScriptsModelImportSupported(resolverCtx.projectGradleVersion)) {
+                    val model = toolingModelsProvider.getProjectModel(it, KotlinDslScriptsModel::class.java)
+                    if (model != null) {
+                        processScriptModel(resolverCtx, model, projectIdentifier)
+                    }
                }
+
+                saveGradleBuildEnvironment(resolverCtx)
            }
        }
    }
@@ -5,6 +5,7 @@

 package org.jetbrains.kotlin.idea.scripting.gradle.importing

+import com.intellij.openapi.components.ServiceManager
 import com.intellij.openapi.externalSystem.model.task.ExternalSystemTaskId
 import com.intellij.openapi.externalSystem.model.task.ExternalSystemTaskNotificationListener
 import com.intellij.openapi.externalSystem.model.task.ExternalSystemTaskNotificationListenerAdapter
@@ -13,6 +14,7 @@ import org.jetbrains.kotlin.idea.core.script.ScriptDefinitionContributor
 import org.jetbrains.kotlin.idea.framework.GRADLE_SYSTEM_ID
 import org.jetbrains.kotlin.idea.scripting.gradle.GradleScriptDefinitionsContributor
 import org.jetbrains.kotlin.idea.scripting.gradle.roots.GradleBuildRootsManager
+import org.jetbrains.plugins.gradle.service.GradleInstallationManager
 import java.util.*

 class KotlinDslSyncListener : ExternalSystemTaskNotificationListenerAdapter() {
@@ -45,8 +47,15 @@ class KotlinDslSyncListener : ExternalSystemTaskNotificationListenerAdapter() {
        // project may be null in case of new project
        val project = id.findProject() ?: return

+        if (sync.gradleHome == null) {
+            sync.gradleHome = ServiceManager
+                .getService(GradleInstallationManager::class.java)
+                .getGradleHome(project, sync.workingDir)
+                ?.canonicalPath
+        }
+
        @Suppress("DEPRECATION")
-        ScriptDefinitionContributor.find<GradleScriptDefinitionsContributor>(project)?.reloadIfNecessary()
+        ScriptDefinitionContributor.find<GradleScriptDefinitionsContributor>(project)?.reloadIfNeeded(sync.workingDir, sync.gradleHome)

        saveScriptModels(project, sync)
    }
@@ -96,12 +96,17 @@ class Imported(
            builder.sdks.addSdk(javaHome)
        }

-        val definitions = GradleScriptDefinitionsContributor.getDefinitions(builder.project)
+        val definitions = GradleScriptDefinitionsContributor.getDefinitions(builder.project, pathPrefix, data.gradleHome)
+        if (definitions == null) {
+            // needed to recreate classRoots if correct script definitions weren't loaded at this moment
+            // in this case classRoots will be recreated after script definitions update
+            builder.useCustomScriptDefinition()
+        }

-        builder.addTemplateClassesRoots(data.templateClasspath)
+        builder.addTemplateClassesRoots(GradleScriptDefinitionsContributor.getDefinitionsTemplateClasspath(data.gradleHome))

        data.models.forEach { script ->
-            val definition = selectScriptDefinition(script, definitions)
+            val definition = definitions?.let { selectScriptDefinition(script, it) }

            builder.addCustom(
                script.file,
@@ -10,6 +10,6 @@ import org.jetbrains.kotlin.idea.scripting.gradle.importing.KotlinDslScriptModel
 data class GradleBuildRootData(
    val importTs: Long,
    val projectRoots: Collection<String>,
-    val templateClasspath: Collection<String>,
+    val gradleHome: String,
    val models: Collection<KotlinDslScriptModel>
 )
@@ -19,7 +19,7 @@ import java.io.DataInputStream
 import java.io.DataOutput

 internal object GradleBuildRootDataSerializer {
-    private val attribute = FileAttribute("kotlin-dsl-script-models", 6, false)
+    private val attribute = FileAttribute("kotlin-dsl-script-models", 7, false)

    fun read(buildRoot: VirtualFile): GradleBuildRootData? {
        return attribute.readAttribute(buildRoot)?.use {
@@ -46,7 +46,7 @@ internal object GradleBuildRootDataSerializer {
 internal fun writeKotlinDslScriptModels(output: DataOutput, data: GradleBuildRootData) {
    val strings = StringsPool.writer(output)
    strings.addStrings(data.projectRoots)
-    strings.addStrings(data.templateClasspath)
+    strings.addString(data.gradleHome)
    data.models.forEach {
        strings.addString(it.file)
        strings.addStrings(it.classPath)
@@ -56,7 +56,7 @@ internal fun writeKotlinDslScriptModels(output: DataOutput, data: GradleBuildRoo
    strings.writeHeader()
    output.writeLong(data.importTs)
    strings.writeStringIds(data.projectRoots)
-    strings.writeStringIds(data.templateClasspath)
+    strings.writeStringId(data.gradleHome)
    output.writeList(data.models) {
        strings.writeStringId(it.file)
        output.writeString(it.inputs.sections)
@@ -72,7 +72,7 @@ internal fun readKotlinDslScriptModels(input: DataInputStream, buildRoot: String

    val importTs = input.readLong()
    val projectRoots = strings.readStrings()
-    val templateClasspath = strings.readStrings()
+    val gradleHome = strings.readString()
    val models = input.readList {
        KotlinDslScriptModel(
            strings.readString(),
@@ -84,7 +84,7 @@ internal fun readKotlinDslScriptModels(input: DataInputStream, buildRoot: String
        )
    }

-    return GradleBuildRootData(importTs, projectRoots, templateClasspath, models)
+    return GradleBuildRootData(importTs, projectRoots, gradleHome, models)
 }

 private object StringsPool {
@@ -19,6 +19,8 @@ abstract class GradleBuildRootsLocator {

    abstract fun getScriptInfo(localPath: String): GradleScriptInfo?

+    fun getAllRoots(): Collection<GradleBuildRoot> = roots.list
+
    fun getBuildRootByWorkingDir(gradleWorkingDir: String) =
        roots.getBuildByRootDir(gradleWorkingDir)

@@ -30,7 +30,9 @@ import org.jetbrains.kotlin.idea.scripting.gradle.roots.GradleBuildRoot.Importin
 import org.jetbrains.kotlin.idea.util.application.runReadAction
 import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.utils.addToStdlib.firstIsInstance
+import org.jetbrains.kotlin.utils.addToStdlib.firstIsInstanceOrNull
 import org.jetbrains.plugins.gradle.config.GradleSettingsListenerAdapter
+import org.jetbrains.plugins.gradle.service.GradleInstallationManager
 import org.jetbrains.plugins.gradle.settings.*
 import org.jetbrains.plugins.gradle.util.GradleConstants
 import java.io.File
@@ -73,7 +75,7 @@ class GradleBuildRootsManager(val project: Project) : GradleBuildRootsLocator(),
            if (value != field) {
                field = value
                roots.list.toList().forEach {
-                    reloadBuildRoot(it.pathPrefix)
+                    reloadBuildRoot(it.pathPrefix, null)
                }
            }
        }
@@ -135,25 +137,31 @@ class GradleBuildRootsManager(val project: Project) : GradleBuildRootsLocator(),
    }

    fun markImportingInProgress(workingDir: String, inProgress: Boolean = true) {
-        actualizeBuildRoot(workingDir)?.importing?.set(if (inProgress) importing else updated)
+        actualizeBuildRoot(workingDir, null)?.importing?.set(if (inProgress) importing else updated)
        updateNotifications { it.startsWith(workingDir) }
    }

    fun update(sync: KotlinDslGradleBuildSync) {
+        val oldRoot = actualizeBuildRoot(sync.workingDir, sync.gradleVersion) ?: return
+
        // fast path for linked gradle builds without .gradle.kts support
        if (sync.models.isEmpty()) {
-            val root = getBuildRootByWorkingDir(sync.workingDir) ?: return
-            if (root is Imported && root.data.models.isEmpty()) return
+            if (oldRoot is Imported && oldRoot.data.models.isEmpty()) return
        }

-        val oldRoot = actualizeBuildRoot(sync.workingDir) ?: return
+        if (oldRoot is Legacy) return
+
        oldRoot.importing.set(updatingCaches)

        try {
-            if (oldRoot is Legacy) return
+            // TODO: can gradleHome be null, what to do in this case
+            val gradleHome = sync.gradleHome
+            if (gradleHome == null) {
+                scriptingInfoLog("Cannot find valid gradle home for ${sync.gradleHome} with version = ${sync.gradleVersion}, script models cannot be saved")
+                return
+            }

-            val templateClasspath = GradleScriptDefinitionsContributor.getDefinitionsTemplateClasspath(project)
-            val newData = GradleBuildRootData(sync.ts, sync.projectRoots, templateClasspath, sync.models)
+            val newData = GradleBuildRootData(sync.ts, sync.projectRoots, gradleHome, sync.models)
            val mergedData = if (sync.failed && oldRoot is Imported) merge(oldRoot.data, newData) else newData

            val lastModifiedFilesReset = LastModifiedFiles()
@@ -175,7 +183,7 @@ class GradleBuildRootsManager(val project: Project) : GradleBuildRootsLocator(),
        val models = old.models.associateByTo(mutableMapOf()) { it.file }
        new.models.associateByTo(models) { it.file }

-        return GradleBuildRootData(new.importTs, roots, new.templateClasspath, models.values)
+        return GradleBuildRootData(new.importTs, roots, new.gradleHome, models.values)
    }

    private val modifiedFilesCheckScheduled = AtomicBoolean()
@@ -206,7 +214,7 @@ class GradleBuildRootsManager(val project: Project) : GradleBuildRootsLocator(),
            // detect gradle version change
            val buildDir = findGradleWrapperPropertiesBuildDir(file)
            if (buildDir != null) {
-                actualizeBuildRoot(buildDir)
+                actualizeBuildRoot(buildDir, null)
            }
        }
    }
@@ -239,11 +247,12 @@ class GradleBuildRootsManager(val project: Project) : GradleBuildRootsLocator(),
            }

            override fun onGradleHomeChange(oldPath: String?, newPath: String?, linkedProjectPath: String) {
-                reloadBuildRoot(linkedProjectPath)
+                val version = GradleInstallationManager.getGradleVersion(newPath)
+                reloadBuildRoot(linkedProjectPath, version)
            }

            override fun onGradleDistributionTypeChange(currentValue: DistributionType?, linkedProjectPath: String) {
-                reloadBuildRoot(linkedProjectPath)
+                reloadBuildRoot(linkedProjectPath, null)
            }
        }

@@ -260,59 +269,69 @@ class GradleBuildRootsManager(val project: Project) : GradleBuildRootsLocator(),
     * Actually this should be true, but we may miss some change events.
     * For that cases we are rechecking this on each Gradle Project sync (importing/reimporting)
     */
-    private fun actualizeBuildRoot(workingDir: String): GradleBuildRoot? {
+    private fun actualizeBuildRoot(workingDir: String, gradleVersion: String?): GradleBuildRoot? {
        val actualSettings = getGradleProjectSettings(workingDir)
        val buildRoot = getBuildRootByWorkingDir(workingDir)

+        val version = gradleVersion ?: actualSettings?.resolveGradleVersion()?.version
        return when {
-            buildRoot != null -> when {
-                !buildRoot.checkActual(actualSettings) -> reloadBuildRoot(workingDir)
-                else -> buildRoot
+            buildRoot != null -> {
+                when {
+                    !buildRoot.checkActual(version) -> reloadBuildRoot(workingDir, version)
+                    else -> buildRoot
+                }
+            }
+            actualSettings != null && version != null -> {
+                loadLinkedRoot(actualSettings, version)
            }
-            actualSettings != null -> loadLinkedRoot(actualSettings)
            else -> null
        }
    }

-    private fun GradleBuildRoot.checkActual(actualSettings: GradleProjectSettings?): Boolean {
-        if (actualSettings == null) return false
+    private fun GradleBuildRoot.checkActual(version: String?): Boolean {
+        if (version == null) return false

        val knownAsSupported = this !is Legacy
-        val shouldBeSupported = kotlinDslScriptsModelImportSupported(actualSettings.resolveGradleVersion().version)
+        val shouldBeSupported = kotlinDslScriptsModelImportSupported(version)
        return knownAsSupported == shouldBeSupported
    }

-    private fun reloadBuildRoot(rootPath: String): GradleBuildRoot? {
+    private fun reloadBuildRoot(rootPath: String, version: String?): GradleBuildRoot? {
        val settings = getGradleProjectSettings(rootPath)
        if (settings == null) {
            remove(rootPath)
            return null
        } else {
-            val newRoot = loadLinkedRoot(settings)
+            val gradleVersion = version ?: settings.resolveGradleVersion().version
+            val newRoot = loadLinkedRoot(settings, gradleVersion)
            add(newRoot)
            return newRoot
        }
    }

-    private fun loadLinkedRoot(settings: GradleProjectSettings): GradleBuildRoot {
+    private fun loadLinkedRoot(settings: GradleProjectSettings, version: String = getGradleVersion(project, settings)): GradleBuildRoot {
        if (!enabled) {
            return Legacy(settings)
        }

-        val supported = kotlinDslScriptsModelImportSupported(settings.resolveGradleVersion().version)
+        val supported = kotlinDslScriptsModelImportSupported(version)

        return when {
-            supported -> tryLoadFromFsCache(settings) ?: New(settings)
+            supported -> tryLoadFromFsCache(settings, version) ?: New(settings)
            else -> Legacy(settings)
        }
    }

-    private fun tryLoadFromFsCache(settings: GradleProjectSettings) =
-        tryCreateImportedRoot(settings.externalProjectPath) {
+    private fun tryLoadFromFsCache(settings: GradleProjectSettings, version: String): Imported? {
+        return tryCreateImportedRoot(settings.externalProjectPath) {
            GradleBuildRootDataSerializer.read(it)?.let { data ->
+                val gradleHome = data.gradleHome
+                if (gradleHome.isNotBlank() && GradleInstallationManager.getGradleVersion(gradleHome) != version) return@let null
+
                addFromSettings(data, settings)
            }
        }
+    }

    private fun addFromSettings(
        data: GradleBuildRootData,
@@ -326,6 +345,7 @@ class GradleBuildRootsManager(val project: Project) : GradleBuildRootsLocator(),
    ): Imported? {
        val buildRoot = VfsUtil.findFile(Paths.get(externalProjectPath), true) ?: return null
        val data = dataProvider(buildRoot) ?: return null
+        // TODO: can be outdated, should be taken from sync
        val javaHome = ExternalSystemApiUtil
            .getExecutionSettings<GradleExecutionSettings>(project, externalProjectPath, GradleConstants.SYSTEM_ID)
            .javaHome?.let { File(it) }
@@ -406,7 +426,7 @@ class GradleBuildRootsManager(val project: Project) : GradleBuildRootsLocator(),

    private fun updateFloatingAction(file: VirtualFile) {
        if (isConfigurationOutOfDate(file)) {
-            scriptConfigurationsNeedToBeUpdated(project)
+            scriptConfigurationsNeedToBeUpdated(project, file)
        } else {
            scriptConfigurationsAreUpToDate(project)
        }
@@ -84,4 +84,8 @@ fun scriptingDebugLog(message: () -> String) {
    if (logger.isDebugEnabled) {
        logger.debug("[KOTLIN_GRADLE_DSL] ${message()}")
    }
+}
+
+fun scriptingInfoLog(message: String) {
+    logger.info("[KOTLIN_GRADLE_DSL] $message")
 }
@@ -22,7 +22,7 @@ class GradleBuildRootDataSerializerTest {
        val data = GradleBuildRootData(
            123,
            listOf("a", "b", "c"),
-            listOf("a"),
+            "a",
            listOf(
                KotlinDslScriptModel(
                    "a",
@@ -39,7 +39,7 @@ open class AbstractGradleBuildRootsLocatorTest {
            GradleBuildRootData(
                ts,
                relativeProjectRoots.map { (pathPrefix + it).removeSuffix("/") },
-                listOf(),
+                "",
                listOf()
            ),
            LastModifiedFiles()
@@ -13,6 +13,9 @@ dependencies {
    compileOnly(intellijCoreDep())
    compileOnly(intellijDep())
    compileOnly(intellijPluginDep("gradle"))
+    Platform[193].orHigher {
+        compileOnly(intellijPluginDep("gradle-java"))
+    }

    testImplementation(projectTests(":idea"))
    testImplementation(project(":libraries:tools:new-project-wizard:new-project-wizard-cli"))
@@ -23,7 +26,7 @@ dependencies {
    testImplementation(intellijDep())
    testImplementation(intellijPluginDep("gradle"))

-
+    testImplementation(projectTests(":idea:idea-gradle"))

    testRuntimeOnly(toolsJar())
    testRuntimeOnly(project(":plugins:kapt3-idea"))
@@ -44,9 +47,9 @@ dependencies {
    }

    Platform[192].orHigher {
-        compileOnly(intellijPluginDep("java")) { includeJars("java-api", "java-impl") }
-        testCompileOnly(intellijPluginDep("java")) { includeJars("java-api", "java-impl") }
-        testRuntimeOnly(intellijPluginDep("java")) { includeJars("java-api") }
+        compileOnly(intellijPluginDep("java"))
+        testCompileOnly(intellijPluginDep("java"))
+        testRuntimeOnly(intellijPluginDep("java"))
    }
 }

@@ -6,11 +6,23 @@
 package org.jetbrains.kotlin.tools.projectWizard.wizard

 import com.intellij.openapi.application.impl.ApplicationInfoImpl
+import com.intellij.openapi.components.service
 import com.intellij.openapi.externalSystem.util.ExternalSystemApiUtil
+import com.intellij.openapi.projectRoots.JavaSdk
+import com.intellij.openapi.projectRoots.JavaSdkType
 import com.intellij.openapi.projectRoots.ProjectJdkTable
 import com.intellij.openapi.projectRoots.SimpleJavaSdkType
+import com.intellij.openapi.projectRoots.impl.SdkConfigurationUtil
+import com.intellij.openapi.vfs.LocalFileSystem
 import com.intellij.testFramework.IdeaTestUtil
 import com.intellij.testFramework.PlatformTestCase
+import org.jetbrains.kotlin.diagnostics.Severity
+import org.jetbrains.kotlin.diagnostics.rendering.DefaultErrorMessages
+import org.jetbrains.kotlin.idea.caches.resolve.analyzeWithContent
+import org.jetbrains.kotlin.idea.codeInsight.gradle.ExternalSystemImportingTestCase
+import org.jetbrains.kotlin.idea.core.script.ScriptConfigurationManager
+import org.jetbrains.kotlin.idea.core.script.configuration.utils.getKtFile
+import org.jetbrains.kotlin.idea.scripting.gradle.getGradleProjectSettings
 import org.jetbrains.kotlin.idea.test.KotlinSdkCreationChecker
 import org.jetbrains.kotlin.idea.test.PluginTestCaseBase
 import org.jetbrains.kotlin.test.testFramework.runWriteAction
@@ -27,6 +39,7 @@ import org.jetbrains.plugins.gradle.settings.GradleProjectSettings
 import org.jetbrains.plugins.gradle.settings.GradleSettings
 import org.jetbrains.plugins.gradle.util.GradleConstants
 import org.jetbrains.plugins.gradle.util.GradleEnvironment
+import java.io.File
 import java.nio.file.Files
 import java.nio.file.Path
 import java.nio.file.Paths
@@ -40,8 +53,9 @@ abstract class AbstractNewWizardProjectImportTest : PlatformTestCase() {
    override fun setUp() {
        super.setUp()
        runWriteAction {
-            val sdk = SimpleJavaSdkType().createJdk(SDK_NAME, IdeaTestUtil.requireRealJdkHome())
-            PluginTestCaseBase.addJdk(testRootDisposable, { sdk })
+            PluginTestCaseBase.addJdk(testRootDisposable) {
+                JavaSdk.getInstance().createJdk(SDK_NAME, IdeaTestUtil.requireRealJdkHome(), false)
+            }
        }
        sdkCreationChecker = KotlinSdkCreationChecker()
    }
@@ -56,6 +70,7 @@ abstract class AbstractNewWizardProjectImportTest : PlatformTestCase() {

    fun doTestGradleKts(directoryPath: String) {
        doTest(directoryPath, BuildSystem.GRADLE_KOTLIN_DSL)
+        checkScriptConfigurationsIfAny()
    }

    fun doTestGradleGroovy(directoryPath: String) {
@@ -77,6 +92,14 @@ abstract class AbstractNewWizardProjectImportTest : PlatformTestCase() {
            prepareGradleBuildSystem(tempDirectory)
        }

+        runWizard(directory, buildSystem, tempDirectory)
+    }
+
+    protected fun runWizard(
+        directory: Path,
+        buildSystem: BuildSystem,
+        tempDirectory: Path
+    ) {
        val wizard = createWizard(directory, buildSystem, tempDirectory)

        val projectDependentServices =
@@ -86,7 +109,10 @@ abstract class AbstractNewWizardProjectImportTest : PlatformTestCase() {
        wizard.apply(projectDependentServices, GenerationPhase.ALL).assertSuccess()
    }

-    protected fun prepareGradleBuildSystem(directory: Path) {
+    protected fun prepareGradleBuildSystem(
+        directory: Path,
+        distributionTypeSettings: DistributionType = DistributionType.WRAPPED
+    ) {
        com.intellij.openapi.components.ServiceManager.getService(project, GradleSettings::class.java)?.apply {
            isOfflineWork = GradleEnvironment.Headless.GRADLE_OFFLINE?.toBoolean() ?: isOfflineWork
            serviceDirectoryPath = GradleEnvironment.Headless.GRADLE_SERVICE_DIRECTORY ?: serviceDirectoryPath
@@ -99,12 +125,36 @@ abstract class AbstractNewWizardProjectImportTest : PlatformTestCase() {
                isUseAutoImport = false
                isUseQualifiedModuleNames = true
                gradleJvm = SDK_NAME
-                distributionType = DistributionType.WRAPPED
+                distributionType = distributionTypeSettings
            }
            ExternalSystemApiUtil.getSettings(project, GradleConstants.SYSTEM_ID).linkProject(settings)
        }
    }

+    protected fun checkScriptConfigurationsIfAny() {
+        if (is192()) return
+
+        val settings = getGradleProjectSettings(project).firstOrNull() ?: error("Cannot find linked gradle project: ${project.basePath}")
+        val scripts = File(settings.externalProjectPath).walkTopDown().filter {
+            it.name.endsWith("gradle.kts")
+        }
+        scripts.forEach {
+            val virtualFile = LocalFileSystem.getInstance().refreshAndFindFileByIoFile(it)!!
+            val psiFile = project.getKtFile(virtualFile) ?: error("Cannot find KtFile for $it")
+            assertTrue(
+                "Configuration for ${it.path} is missing",
+                project.service<ScriptConfigurationManager>().hasConfiguration(psiFile)
+            )
+
+            val bindingContext = psiFile.analyzeWithContent()
+
+            val diagnostics = bindingContext.diagnostics.filter { it.severity == Severity.ERROR }
+            assert(diagnostics.isEmpty()) {
+                "Diagnostics list should be empty:\n ${diagnostics.joinToString("\n") { DefaultErrorMessages.render(it) }}"
+            }
+        }
+    }
+
    private fun is192() =
        ApplicationInfoImpl.getShadowInstance().minorVersionMainPart == "2"
                && ApplicationInfoImpl.getShadowInstance().majorVersion == "2019"
@@ -0,0 +1,41 @@
+/*
+ * Copyright 2010-2020 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.tools.projectWizard.wizard
+
+import org.jetbrains.kotlin.tools.projectWizard.cli.BuildSystem
+import org.jetbrains.plugins.gradle.settings.DistributionType
+import org.junit.Test
+import java.nio.file.Files
+import java.nio.file.Paths
+
+class ScriptHighlightingGradleDistributionTypeTest : AbstractProjectTemplateNewWizardProjectImportTest() {
+
+    @Test
+    fun testScriptHighlightingGradleWrapped() {
+        doTest(DistributionType.WRAPPED)
+    }
+
+    @Test
+    fun testScriptHighlightingGradleDefaultWrapped() {
+        doTest(DistributionType.DEFAULT_WRAPPED)
+    }
+
+    @Test
+    fun testScriptHighlightingGradleBundled() {
+        doTest(DistributionType.BUNDLED)
+    }
+
+    private fun doTest(distributionType: DistributionType) {
+        val directory = Paths.get("backendApplication")
+        val tempDirectory = Files.createTempDirectory(null)
+
+        prepareGradleBuildSystem(tempDirectory, distributionType)
+
+        runWizard(directory, BuildSystem.GRADLE_KOTLIN_DSL, tempDirectory)
+
+        checkScriptConfigurationsIfAny()
+    }
+}