From e395a32734888c2ff667ee9e35fdb4b63f24c6c6 Mon Sep 17 00:00:00 2001
From: Natalia Selezneva <natalia.selezneva@jetbrains.com>
Date: Wed, 11 Jul 2018 14:23:18 +0300
Subject: [PATCH] Add quickfix to adding kotlin-script-runtime.jar to classpath

^KT-25402 Fixed
---
 .../kotlin/idea/KotlinBundle.properties       |  1 +
 .../idea/inspections/AddKotlinLibQuickFix.kt  | 24 +++++++++++++++++++
 .../idea/quickfix/JvmQuickFixRegistrar.kt     |  3 +++
 3 files changed, 28 insertions(+)

diff --git a/idea/idea-analysis/src/org/jetbrains/kotlin/idea/KotlinBundle.properties b/idea/idea-analysis/src/org/jetbrains/kotlin/idea/KotlinBundle.properties
index 5451ec5c723..879f759fcf1 100644
--- a/idea/idea-analysis/src/org/jetbrains/kotlin/idea/KotlinBundle.properties
+++ b/idea/idea-analysis/src/org/jetbrains/kotlin/idea/KotlinBundle.properties
@@ -215,6 +215,7 @@ suppress.warning.for=Suppress ''{0}'' for {1} {2}
 reflection.not.found=Reflection not found in the classpath
 add.reflection.to.classpath=Add kotlin-reflect.jar to the classpath
 add.test.to.classpath=Add kotlin-test.jar to the classpath
+add.script.runtime.to.classpath=Add kotlin-script-runtime.jar to the classpath
 
 # Kotlin Compiler Settings Tab
 
diff --git a/idea/idea-jvm/src/org/jetbrains/kotlin/idea/inspections/AddKotlinLibQuickFix.kt b/idea/idea-jvm/src/org/jetbrains/kotlin/idea/inspections/AddKotlinLibQuickFix.kt
index 9c5aeb61268..1a91f0eb31c 100644
--- a/idea/idea-jvm/src/org/jetbrains/kotlin/idea/inspections/AddKotlinLibQuickFix.kt
+++ b/idea/idea-jvm/src/org/jetbrains/kotlin/idea/inspections/AddKotlinLibQuickFix.kt
@@ -37,6 +37,7 @@ import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.psi.KtImportDirective
+import org.jetbrains.kotlin.psi.KtScript
 
 class AddReflectionQuickFix(element: KtElement)
         : AddKotlinLibQuickFix(element, listOf(LibraryJarDescriptor.REFLECT_JAR,
@@ -52,6 +53,29 @@ class AddReflectionQuickFix(element: KtElement)
     }
 }
 
+class AddScriptRuntimeQuickFix(element: KtElement) : AddKotlinLibQuickFix(element, listOf(LibraryJarDescriptor.SCRIPT_RUNTIME_JAR)) {
+    override fun getText() = KotlinBundle.message("add.script.runtime.to.classpath")
+    override fun getFamilyName() = text
+
+    override fun getLibraryDescriptor(module: Module) = MavenExternalLibraryDescriptor(
+        "org.jetbrains.kotlin", "kotlin-script-runtime",
+        getRuntimeLibraryVersion(module) ?: bundledRuntimeVersion()
+    )
+
+    companion object : KotlinSingleIntentionActionFactory() {
+        override fun createAction(diagnostic: Diagnostic): KotlinQuickFixAction<KtElement>? {
+            val ktScript = Errors.MISSING_SCRIPT_BASE_CLASS.cast(diagnostic).psiElement as? KtScript ?: return null
+            val templateClassName = ktScript.kotlinScriptDefinition.value.template.qualifiedName ?: return null
+
+            if (templateClassName.startsWith("kotlin.script.templates.standard")) {
+                return diagnostic.createIntentionForFirstParentOfType(::AddScriptRuntimeQuickFix)
+            }
+
+            return null
+        }
+    }
+}
+
 class AddTestLibQuickFix(element: KtElement)
         : AddKotlinLibQuickFix(element, listOf(LibraryJarDescriptor.TEST_JAR,
                                                LibraryJarDescriptor.TEST_SRC_JAR)) {
diff --git a/idea/idea-jvm/src/org/jetbrains/kotlin/idea/quickfix/JvmQuickFixRegistrar.kt b/idea/idea-jvm/src/org/jetbrains/kotlin/idea/quickfix/JvmQuickFixRegistrar.kt
index faa863a7fa0..02137fc932d 100644
--- a/idea/idea-jvm/src/org/jetbrains/kotlin/idea/quickfix/JvmQuickFixRegistrar.kt
+++ b/idea/idea-jvm/src/org/jetbrains/kotlin/idea/quickfix/JvmQuickFixRegistrar.kt
@@ -19,6 +19,7 @@ package org.jetbrains.kotlin.idea.quickfix
 import org.jetbrains.kotlin.diagnostics.DiagnosticFactory
 import org.jetbrains.kotlin.diagnostics.Errors.*
 import org.jetbrains.kotlin.idea.inspections.AddReflectionQuickFix
+import org.jetbrains.kotlin.idea.inspections.AddScriptRuntimeQuickFix
 import org.jetbrains.kotlin.idea.inspections.AddTestLibQuickFix
 import org.jetbrains.kotlin.resolve.jvm.diagnostics.ErrorsJvm.NO_REFLECTION_IN_CLASS_PATH
 
@@ -36,5 +37,7 @@ class JvmQuickFixRegistrar : QuickFixContributor {
         EXPERIMENTAL_FEATURE_WARNING.registerFactory(ChangeCoroutineSupportFix)
 
         NO_REFLECTION_IN_CLASS_PATH.registerFactory(AddReflectionQuickFix)
+
+        MISSING_SCRIPT_BASE_CLASS.registerFactory(AddScriptRuntimeQuickFix)
     }
 }
\ No newline at end of file