feat: auto import

src/main/java/ee/carlrobert/codegpt/CodeGPTKeys.java

@@ -45,7 +45,6 @@ public class CodeGPTKeys {
       Key.create("codegpt.lastCompletionResponseId");
   public static final Key<ToolWindowEditorFileDetails> TOOLWINDOW_EDITOR_FILE_DETAILS =
       Key.create("proxyai.toolwindowEditorFileDetails");
-
   public static final Key<NextEditDiffViewer> EDITOR_PREDICTION_DIFF_VIEWER =
       Key.create("codegpt.editorPredictionDiffViewer");
 }

src/main/kotlin/ee/carlrobert/codegpt/CodeGPTLookupListener.kt

@@ -9,9 +9,18 @@ import com.intellij.codeInsight.lookup.LookupManagerListener
 import com.intellij.codeInsight.lookup.impl.LookupImpl
 import com.intellij.openapi.components.service
 import com.intellij.openapi.editor.EditorKind
+import com.intellij.openapi.util.TextRange
 import ee.carlrobert.codegpt.codecompletions.CodeCompletionService
+import ee.carlrobert.codegpt.nextedit.NextEditCoordinator
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.SupervisorJob
+import kotlinx.coroutines.launch
 
 class CodeGPTLookupListener : LookupManagerListener {
+
+    private val coroutineScope = CoroutineScope(SupervisorJob() + Dispatchers.Default)
+
     override fun activeLookupChanged(oldLookup: Lookup?, newLookup: Lookup?) {
         if (newLookup is LookupImpl) {
             newLookup.addLookupListener(object : LookupListener {
@@ -26,10 +35,19 @@ class CodeGPTLookupListener : LookupManagerListener {
                     ) {
                         return
                     }
+                    val offset = editor.caretModel.offset
+                    val lineNumber = editor.document.getLineNumber(offset)
+                    val lineEndOffset = editor.document.getLineEndOffset(lineNumber)
 
-                    InlineCompletion.getHandlerOrNull(editor)?.invokeEvent(
-                        InlineCompletionEvent.DirectCall(editor, editor.caretModel.currentCaret)
-                    )
+                    if (editor.document.getText(TextRange(offset, lineEndOffset)).isEmpty()) {
+                        InlineCompletion.getHandlerOrNull(editor)?.invokeEvent(
+                            InlineCompletionEvent.DirectCall(editor, editor.caretModel.currentCaret)
+                        )
+                    } else {
+                        coroutineScope.launch {
+                            NextEditCoordinator.requestNextEdit(editor, editor.document.text, offset)
+                        }
+                    }
                 }
             })
         }

src/main/kotlin/ee/carlrobert/codegpt/autoimport/AutoImportOrchestrator.kt

@@ -0,0 +1,55 @@
+package ee.carlrobert.codegpt.autoimport
+
+import com.intellij.openapi.application.runReadAction
+import com.intellij.openapi.command.WriteCommandAction
+import com.intellij.openapi.diagnostic.thisLogger
+import com.intellij.openapi.editor.Editor
+import com.intellij.openapi.util.TextRange
+import com.intellij.psi.PsiDocumentManager
+import com.intellij.psi.PsiFile
+import com.intellij.psi.PsiReference
+
+data class UnresolvedSymbol(
+    val name: String,
+    val reference: PsiReference,
+    val range: TextRange,
+)
+
+data class ImportCandidate(
+    val fqn: String,
+)
+
+object AutoImportOrchestrator {
+
+    private val logger = thisLogger()
+
+    private val resolvers: List<AutoImportResolver> = listOf(
+        JavaResolver(),
+        KtResolver(),
+    )
+
+    /**
+     * Preview imports by cloning the current editor file's PSI, finding unresolved imports within [range],
+     * applying the best candidates to the clone, and returning the resulting content and list of added import FQNs.
+     * Note: This does not modify the original editor file.
+     */
+    fun previewImports(editor: Editor, range: TextRange? = null): String? {
+        val project = editor.project ?: return null
+        val psiFile =
+            PsiDocumentManager.getInstance(project).getPsiFile(editor.document) ?: return null
+        val clonedPsiFile = psiFile.copy() as? PsiFile ?: return null
+        val rangeToUse = range ?: runReadAction { TextRange(0, editor.document.textLength) }
+        val resolver = resolvers.firstOrNull { it.supports(clonedPsiFile) } ?: return null
+
+        runReadAction { resolver.getUnresolvedImports(psiFile, rangeToUse) }
+            .forEach {
+                WriteCommandAction.runWriteCommandAction(project) {
+                    if (!resolver.applyImport(clonedPsiFile, it)) {
+                        logger.warn("Failed to apply import: $it")
+                    }
+                }
+            }
+
+        return runReadAction { clonedPsiFile.text }
+    }
+}

src/main/kotlin/ee/carlrobert/codegpt/autoimport/AutoImportResolver.kt

@@ -0,0 +1,10 @@
+package ee.carlrobert.codegpt.autoimport
+
+import com.intellij.openapi.util.TextRange
+import com.intellij.psi.PsiFile
+
+interface AutoImportResolver {
+    fun supports(file: PsiFile): Boolean
+    fun getUnresolvedImports(file: PsiFile, searchRange: TextRange): List<String>
+    fun applyImport(file: PsiFile, importFqn: String): Boolean
+}

src/main/kotlin/ee/carlrobert/codegpt/autoimport/JavaResolver.kt

@@ -0,0 +1,116 @@
+package ee.carlrobert.codegpt.autoimport
+
+import com.intellij.openapi.application.runReadAction
+import com.intellij.openapi.components.service
+import com.intellij.openapi.util.TextRange
+import com.intellij.psi.*
+import com.intellij.psi.codeStyle.CodeStyleManager
+import com.intellij.psi.codeStyle.JavaCodeStyleManager
+import com.intellij.psi.search.GlobalSearchScope
+import com.intellij.psi.search.PsiShortNamesCache
+
+internal class JavaResolver : AutoImportResolver {
+    override fun supports(file: PsiFile): Boolean = file is PsiJavaFile
+
+    override fun getUnresolvedImports(
+        file: PsiFile,
+        searchRange: TextRange
+    ): List<String> {
+        val result = mutableListOf<String>()
+        file.accept(object : JavaRecursiveElementWalkingVisitor() {
+            override fun visitReferenceElement(referenceElement: PsiJavaCodeReferenceElement) {
+                super.visitReferenceElement(referenceElement)
+
+                if (!isInRange(referenceElement as PsiReference, searchRange)) return
+
+                val resolved = runReadAction { referenceElement.resolve() }
+                if (resolved == null) {
+                    val name = referenceElement.referenceName ?: referenceElement.canonicalText
+                    if (name.isNotBlank()) {
+                        val range = referenceElement.textRange ?: TextRange.EMPTY_RANGE
+                        val symbol = UnresolvedSymbol(name, referenceElement as PsiReference, range)
+                        bestCandidateFor(symbol, file)?.let {
+                            result.add(it)
+                        }
+                    }
+                }
+            }
+        })
+        return result.distinctBy { it }
+    }
+
+    override fun applyImport(file: PsiFile, importFqn: String): Boolean {
+        if (file !is PsiJavaFile) return false
+
+        val project = file.project
+        val cls = project.service<JavaPsiFacade>().findClass(importFqn, file.resolveScope)
+            ?: return false
+        val javaCodeStyleManager = project.service<JavaCodeStyleManager>()
+        val added = javaCodeStyleManager.addImport(file, cls)
+        if (added) {
+            val importList = file.importList
+            if (importList != null) {
+                val endOffset = importList.nextSibling?.textRange?.startOffset?.let { it + 1 }
+                    ?: importList.textRange.endOffset
+
+                CodeStyleManager.getInstance(project).reformatRange(file,
+                    importList.textRange.startOffset,
+                    endOffset.coerceAtMost(file.textRange.endOffset))
+            }
+            javaCodeStyleManager.optimizeImports(file)
+        }
+        return added
+    }
+
+    private fun bestCandidateFor(symbol: UnresolvedSymbol, file: PsiFile): String? {
+        if (file !is PsiJavaFile) return null
+        val project = file.project
+        val scope = GlobalSearchScope.allScope(project)
+        val classes = runReadAction {
+            PsiShortNamesCache.getInstance(project).getClassesByName(symbol.name, scope)
+        }
+
+        val alreadyImported: Set<String> = buildSet {
+            val list = runReadAction { file.importList }
+            list?.allImportStatements?.forEach { stmt ->
+                val qn = runReadAction { stmt.importReference?.qualifiedName }
+                if (qn != null) add(qn)
+            }
+        }
+        val currentPackage = runReadAction { file.packageName }
+
+        val fqns = classes.mapNotNull { runReadAction { it.qualifiedName } }
+            .filter { qn ->
+                val pkg = qn.substringBeforeLast('.', "")
+                pkg != currentPackage && qn !in alreadyImported
+            }
+            .distinct()
+            .sorted()
+
+        return fqns.asSequence()
+            .filterNot { it.startsWith("com.sun.") || it.startsWith("sun.") }
+            .sortedByDescending { rankImport(it) }
+            .firstOrNull()
+    }
+
+    private fun rankImport(fqn: String): Int {
+        var score = 0
+
+        when {
+            fqn.startsWith("java.util.") -> score += 100
+            fqn.startsWith("java.lang.") -> score += 90
+            fqn.startsWith("java.io.") -> score += 80
+            fqn.startsWith("java.nio.") -> score += 80
+            fqn.startsWith("java.time.") -> score += 70
+            fqn.startsWith("java.awt.") -> score -= 50
+            fqn.startsWith("javax.swing.") -> score -= 30
+        }
+
+        return score
+    }
+
+    private fun isInRange(ref: PsiReference, searchRange: TextRange): Boolean {
+        val refRange = runReadAction { ref.element.textRange } ?: return false
+        return refRange.intersects(searchRange)
+    }
+}

src/main/kotlin/ee/carlrobert/codegpt/autoimport/KtResolver.kt

@@ -0,0 +1,110 @@
+package ee.carlrobert.codegpt.autoimport
+
+import com.intellij.openapi.components.service
+import com.intellij.openapi.util.TextRange
+import com.intellij.psi.JavaPsiFacade
+import com.intellij.psi.PsiElement
+import com.intellij.psi.PsiFile
+import com.intellij.psi.PsiReference
+import com.intellij.psi.codeStyle.CodeStyleManager
+import com.intellij.psi.codeStyle.JavaCodeStyleManager
+import com.intellij.psi.search.GlobalSearchScope
+import com.intellij.psi.search.PsiShortNamesCache
+import org.jetbrains.kotlin.name.FqName
+import org.jetbrains.kotlin.psi.KtFile
+import org.jetbrains.kotlin.psi.KtPsiFactory
+import org.jetbrains.kotlin.psi.KtReferenceExpression
+import org.jetbrains.kotlin.psi.KtTreeVisitorVoid
+import org.jetbrains.kotlin.resolve.ImportPath
+
+internal class KtResolver : AutoImportResolver {
+    override fun supports(file: PsiFile): Boolean = file is KtFile
+
+    override fun getUnresolvedImports(
+        file: PsiFile,
+        searchRange: TextRange
+    ): List<String> {
+        val result = mutableListOf<String>()
+        file.accept(object : KtTreeVisitorVoid() {
+            override fun visitElement(element: PsiElement) {
+                super.visitElement(element)
+
+                if (element is KtReferenceExpression) {
+                    val references = element.references
+                    if (references.isEmpty()) return
+                    references.forEach { ref ->
+                        if (!isInRange(ref, searchRange)) return@forEach
+
+                        if (ref.resolve() == null) {
+                            val name = ref.canonicalText.substringAfterLast('.')
+                            if (name.isNotBlank()) {
+                                val range = element.textRange ?: TextRange.EMPTY_RANGE
+                                val symbol = UnresolvedSymbol(name, ref, range)
+                                bestCandidateFor(symbol, file)?.let {
+                                    result.add(it)
+                                }
+                            }
+                        }
+                    }
+                }
+            }
+        })
+        return result.distinctBy { it }
+    }
+
+    override fun applyImport(file: PsiFile, importFqn: String): Boolean {
+        if (file !is KtFile) return false
+
+        val importList = file.importList ?: return false
+        if (file.importDirectives.any { !it.isAllUnder && it.importedFqName?.asString() == importFqn }) return false
+
+        val project = file.project
+        val directive = KtPsiFactory(project).createImportDirective(
+            ImportPath(FqName(importFqn), false)
+        )
+        importList.add(directive)
+
+        val endOffset = importList.nextSibling?.textRange?.startOffset?.let { it + 1 }
+            ?: importList.textRange.endOffset
+
+        CodeStyleManager.getInstance(project).reformatRange(
+            file,
+            importList.textRange.startOffset,
+            endOffset.coerceAtMost(file.textRange.endOffset)
+        )
+
+        return true
+    }
+
+    private fun bestCandidateFor(symbol: UnresolvedSymbol, file: PsiFile): String? {
+        if (file !is KtFile) return null
+        val project = file.project
+        val facade = JavaPsiFacade.getInstance(project)
+        val scope = GlobalSearchScope.allScope(project)
+
+        val name = symbol.name
+        val classes = facade.findClasses(name, scope)
+        val alreadyImported = file.importDirectives
+            .mapNotNull { if (!it.isAllUnder) it.importedFqName?.asString() else null }
+            .toSet()
+
+        val fromJavaFacade = classes.mapNotNull { it.qualifiedName }
+
+        val fromShortNamesCache = PsiShortNamesCache.getInstance(project)
+            .getClassesByName(name, scope)
+            .mapNotNull { it.qualifiedName }
+
+        val fqns = (fromJavaFacade + fromShortNamesCache)
+            .filter { qn -> qn !in alreadyImported }
+            .distinct()
+            .sorted()
+
+        return fqns.asSequence().firstOrNull()
+    }
+
+    private fun isInRange(ref: PsiReference, searchRange: TextRange?): Boolean {
+        searchRange ?: return true
+        val range = ref.element.textRange ?: return false
+        return range.intersects(searchRange)
+    }
+}

src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionInsertHandler.kt

@@ -3,37 +3,55 @@ package ee.carlrobert.codegpt.codecompletions
 import com.intellij.codeInsight.inline.completion.InlineCompletionInsertEnvironment
 import com.intellij.codeInsight.inline.completion.InlineCompletionInsertHandler
 import com.intellij.codeInsight.inline.completion.elements.InlineCompletionElement
+import com.intellij.openapi.application.EDT
 import com.intellij.openapi.application.runReadAction
 import com.intellij.openapi.components.service
 import com.intellij.openapi.editor.Editor
 import com.intellij.openapi.util.TextRange
 import ee.carlrobert.codegpt.CodeGPTKeys
+import ee.carlrobert.codegpt.autoimport.AutoImportOrchestrator
 import ee.carlrobert.codegpt.codecompletions.edit.GrpcClientService
 import ee.carlrobert.codegpt.nextedit.NextEditCoordinator
-import kotlinx.coroutines.CoroutineScope
-import kotlinx.coroutines.Dispatchers
-import kotlinx.coroutines.SupervisorJob
-import kotlinx.coroutines.launch
+import ee.carlrobert.codegpt.nextedit.NextEditDiffViewer
+import ee.carlrobert.service.NextEditResponse
+import kotlinx.coroutines.*
 
 class CodeCompletionInsertHandler : InlineCompletionInsertHandler {
 
+    private val coroutineScope = CoroutineScope(SupervisorJob() + Dispatchers.Default)
+
     override fun afterInsertion(
         environment: InlineCompletionInsertEnvironment,
         elements: List<InlineCompletionElement>
     ) {
-        val editor = environment.editor
-        val completion = elements.first().text
-        acceptCompletion(completion, editor)
+        coroutineScope.launch {
+            val completion = elements.first().text
+            val editor = environment.editor
+            acceptCompletion(completion, editor)
 
-        val caretOffset = runReadAction { editor.caretModel.offset }
-        requestNextEditAsync(editor, caretOffset)
+            val currentContent = runReadAction { editor.document.text }
+            val caretOffset = runReadAction { editor.caretModel.offset }
+
+            val completionRange = runReadAction { editor.getUserData(CodeGPTKeys.RECENT_COMPLETION_RANGE) }
+            val contentWithImports = AutoImportOrchestrator.previewImports(editor, completionRange)
+            if (contentWithImports == null) {
+                NextEditCoordinator.requestNextEdit(editor, currentContent, caretOffset, false)
+            } else {
+                withContext(Dispatchers.EDT) {
+                    if (contentWithImports.isNotEmpty() && contentWithImports != currentContent) {
+                        showImportDiffViewer(editor, currentContent, contentWithImports)
+                    }
+                }
+            }
+        }
     }
 
     private fun acceptCompletion(completion: String, editor: Editor) {
-        val caret = runReadAction { editor.caretModel.offset }
-        val start = (caret - completion.length).coerceAtLeast(0)
+        val caretOffset = runReadAction { editor.caretModel.offset }
+        val startLine = runReadAction { editor.document.getLineNumber(caretOffset - completion.length) }
+        val start = runReadAction { editor.document.getLineStartOffset(startLine) }
         CodeGPTKeys.RECENT_COMPLETION_TEXT.set(editor, completion)
-        CodeGPTKeys.RECENT_COMPLETION_RANGE.set(editor, TextRange(start, caret))
+        CodeGPTKeys.RECENT_COMPLETION_RANGE.set(editor, TextRange(start, caretOffset))
 
         val responseId = CodeGPTKeys.LAST_COMPLETION_RESPONSE_ID.get(editor)
         if (responseId != null) {
@@ -43,14 +61,17 @@ class CodeCompletionInsertHandler : InlineCompletionInsertHandler {
         }
     }
 
-    private fun requestNextEditAsync(editor: Editor, caretOffset: Int) {
-        CoroutineScope(SupervisorJob() + Dispatchers.IO).launch {
-            NextEditCoordinator.requestNextEdit(
-                editor,
-                editor.document.text,
-                caretOffset,
-                false
-            )
-        }
+    private fun showImportDiffViewer(
+        editor: Editor,
+        oldRevision: String,
+        contentWithImports: String
+    ) {
+        val importResponse = NextEditResponse.newBuilder()
+            .setId("import-${System.currentTimeMillis()}")
+            .setOldRevision(oldRevision)
+            .setNextRevision(contentWithImports)
+            .build()
+
+        NextEditDiffViewer.displayNextEdit(editor, importResponse)
     }
 }

src/main/kotlin/ee/carlrobert/codegpt/nextedit/NextEditDiffViewer.kt

@@ -10,6 +10,7 @@ import com.intellij.diff.tools.fragmented.UnifiedDiffViewer
 import com.intellij.diff.util.DiffUtil
 import com.intellij.diff.util.Side
 import com.intellij.openapi.Disposable
+import com.intellij.openapi.application.EDT
 import com.intellij.openapi.application.runReadAction
 import com.intellij.openapi.components.service
 import com.intellij.openapi.editor.Document
@@ -26,10 +27,10 @@ import com.intellij.testFramework.LightVirtualFile
 import com.intellij.util.application
 import com.intellij.util.concurrency.annotations.RequiresEdt
 import com.intellij.util.ui.JBUI
-import ee.carlrobert.codegpt.AutoImportHandler
 import ee.carlrobert.codegpt.CodeGPTKeys
 import ee.carlrobert.codegpt.codecompletions.edit.GrpcClientService
 import ee.carlrobert.service.NextEditResponse
+import kotlinx.coroutines.*
 import java.awt.Dimension
 import java.awt.Point
 import javax.swing.JComponent
@@ -49,6 +50,7 @@ class NextEditDiffViewer(
     private val visibleAreaListener: VisibleAreaListener
     private val caretListener: CaretListener
     private val grpcService = project?.service<GrpcClientService>()
+    private val coroutineScope = CoroutineScope(SupervisorJob() + Dispatchers.Default)
 
     private var applyInProgress = false
 
@@ -121,8 +123,6 @@ class NextEditDiffViewer(
             scheduleRediff()
         }
 
-        AutoImportHandler.handleAutoImports(mainEditor)
-
         application.executeOnPooledThread {
             val cursor = runReadAction { mainEditor.caretModel.offset }
             grpcService?.acceptEdit(nextEditResponse.id, leftText, rightText, cursor)
@@ -162,7 +162,12 @@ class NextEditDiffViewer(
         val length = document.textLength
         val safeStart = maxOf(0, minOf(start, length))
         val safeEnd = maxOf(safeStart, minOf(end, length))
-        return if (safeStart < length && safeEnd > safeStart) document.getText(TextRange(safeStart, safeEnd)) else ""
+        return if (safeStart < length && safeEnd > safeStart) document.getText(
+            TextRange(
+                safeStart,
+                safeEnd
+            )
+        ) else ""
     }
 
     private fun getClosestChange(): UnifiedDiffChange? {
@@ -176,8 +181,10 @@ class NextEditDiffViewer(
                 val rightStart = change.lineFragment.startOffset2
                 val rightEnd = change.lineFragment.endOffset2
 
-                val validLeft = leftStart >= 0 && leftEnd >= leftStart && leftStart < leftDoc.textLength
-                val validRight = rightStart >= 0 && rightEnd >= rightStart && rightStart < rightDoc.textLength
+                val validLeft =
+                    leftStart >= 0 && leftEnd >= leftStart && leftStart < leftDoc.textLength
+                val validRight =
+                    rightStart >= 0 && rightEnd >= rightStart && rightStart < rightDoc.textLength
 
                 if (!validLeft || !validRight) return@filter false
 
@@ -212,16 +219,19 @@ class NextEditDiffViewer(
             val change = getClosestChange() ?: return
             if (popup.isDisposed) return
 
-            adjustPopupSize(popup, myEditor)
+            coroutineScope.launch {
+                withContext(Dispatchers.EDT) {
+                    adjustPopupSize(popup, myEditor)
+                    val adjustedLocation = getAdjustedPopupLocation(
+                        mainEditor,
+                        change.lineFragment.startOffset1,
+                        popup.size
+                    )
 
-            val adjustedLocation = getAdjustedPopupLocation(
-                mainEditor,
-                change.lineFragment.startOffset1,
-                popup.size
-            )
-
-            if (popup.isVisible && !popup.isDisposed) {
-                popup.setLocation(adjustedLocation)
+                    if (popup.isVisible && !popup.isDisposed) {
+                        popup.setLocation(adjustedLocation)
+                    }
+                }
             }
         }
     }
@@ -247,8 +257,16 @@ class NextEditDiffViewer(
     }
 
     private fun computeCompactSize(change: UnifiedDiffChange): Dimension {
-        val leftText = safeGetText(getDocument(Side.LEFT), change.lineFragment.startOffset1, change.lineFragment.endOffset1)
-        val rightText = safeGetText(getDocument(Side.RIGHT), change.lineFragment.startOffset2, change.lineFragment.endOffset2)
+        val leftText = safeGetText(
+            getDocument(Side.LEFT),
+            change.lineFragment.startOffset1,
+            change.lineFragment.endOffset1
+        )
+        val rightText = safeGetText(
+            getDocument(Side.RIGHT),
+            change.lineFragment.startOffset2,
+            change.lineFragment.endOffset2
+        )
 
         fun linesOf(s: String): List<String> {
             val cleaned = s.replace("\r", "")

src/main/kotlin/ee/carlrobert/codegpt/settings/configuration/ConfigurationSettings.kt

@@ -59,8 +59,7 @@ class ChatCompletionSettingsState : BaseState() {
 class CodeCompletionSettingsState : BaseState() {
     var treeSitterProcessingEnabled by property(true)
     var gitDiffEnabled by property(true)
-    var collectDependencyStructure by property(true)
+    var collectDependencyStructure by property(false)
     var contextAwareEnabled by property(false)
     var psiStructureAnalyzeDepth by property(2)
-    var myAwesomeFeatureEnabled by property(true)
 }

src/main/kotlin/ee/carlrobert/codegpt/settings/models/ModelRegistry.kt

@@ -351,7 +351,7 @@ class ModelRegistry {
     private fun getNextEditModels(): List<ModelSelection> {
         return listOf(
             ModelSelection(ServiceType.PROXYAI, MERCURY_CODER, "Mercury Coder"),
-            ModelSelection(ServiceType.INCEPTION, MERCURY_CODER, "Mercury Coder")
+            ModelSelection(ServiceType.INCEPTION, MERCURY_CODER_NES_PREVIEW, "Mercury Coder (NES Preview)")
         )
     }
 
@@ -660,6 +660,7 @@ class ModelRegistry {
         const val LLAMA_3_2_3B_INSTRUCT = "llama-3.2-3b-instruct"
 
         const val MERCURY_CODER = "mercury-coder"
+        const val MERCURY_CODER_NES_PREVIEW = "mercury-coder-nes-preview"
 
         @JvmStatic
         fun getInstance(): ModelRegistry {

src/main/resources/META-INF/plugin-cpp.xml

@@ -0,0 +1,4 @@
+<idea-plugin>
+  <extensions defaultExtensionNs="com.intellij">
+  </extensions>
+</idea-plugin>

src/main/resources/META-INF/plugin-java.xml

@@ -1,11 +1,11 @@
 <idea-plugin>
   <projectListeners>
     <listener topic="com.intellij.openapi.compiler.CompilationStatusListener"
-      class="ee.carlrobert.codegpt.ProjectCompilationStatusListener" />
+      class="ee.carlrobert.codegpt.ProjectCompilationStatusListener"/>
   </projectListeners>
 
   <extensions defaultExtensionNs="com.intellij">
     <applicationService
       serviceImplementation="ee.carlrobert.codegpt.codecompletions.psi.JavaContextFinder"/>
   </extensions>
-</idea-plugin>
+</idea-plugin>

src/main/resources/META-INF/plugin-javascript.xml

@@ -0,0 +1,4 @@
+<idea-plugin>
+  <extensions defaultExtensionNs="com.intellij">
+  </extensions>
+</idea-plugin>

src/main/resources/META-INF/plugin-kotlin.xml

@@ -1,6 +1,6 @@
 <idea-plugin>
-    <extensions defaultExtensionNs="com.intellij">
-        <applicationService
-                serviceImplementation="ee.carlrobert.codegpt.psistructure.KotlinFileAnalyzer"/>
-    </extensions>
+  <extensions defaultExtensionNs="com.intellij">
+    <applicationService
+      serviceImplementation="ee.carlrobert.codegpt.psistructure.KotlinFileAnalyzer"/>
+  </extensions>
 </idea-plugin>

src/main/resources/META-INF/plugin-python.xml

@@ -3,4 +3,4 @@
     <applicationService
       serviceImplementation="ee.carlrobert.codegpt.codecompletions.psi.PythonContextFinder"/>
   </extensions>
-</idea-plugin>
+</idea-plugin>

src/main/resources/META-INF/plugin.xml

@@ -7,6 +7,8 @@
     <depends optional="true" config-file="plugin-kotlin.xml">org.jetbrains.kotlin</depends>
     <depends optional="true" config-file="plugin-java.xml">com.intellij.modules.java</depends>
     <depends optional="true" config-file="plugin-python.xml">com.intellij.modules.python</depends>
+    <depends optional="true" config-file="plugin-javascript.xml">JavaScript</depends>
+    <depends optional="true" config-file="plugin-cpp.xml">com.intellij.modules.cpp</depends>
     <!--    <depends optional="true" config-file="plugin-go.xml">org.jetbrains.plugins.go</depends>-->
     <!--    <depends optional="true" config-file="plugin-ruby.xml">com.intellij.modules.ruby</depends>-->
     <!--    <depends optional="true" config-file="plugin-php.xml">com.jetbrains.php</depends>-->

src/test/kotlin/ee/carlrobert/codegpt/autoimport/AutoImportOrchestratorTest.kt

@@ -1,4 +1,529 @@
 package ee.carlrobert.codegpt.autoimport
 
-class AutoImportOrchestratorTest {
-}
+import com.intellij.openapi.util.TextRange
+import org.assertj.core.api.Assertions.assertThat
+import testsupport.IntegrationTest
+
+class AutoImportOrchestratorTest : IntegrationTest() {
+
+    fun testPreviewImportsForJavaFile() {
+        myFixture.addFileToProject(
+            "com/test/util/ListClass.java",
+            "package com.test.util; public class ListClass<E> {}"
+        )
+        myFixture.addFileToProject(
+            "com/test/lang/StringHelper.java",
+            "package com.test.lang; public class StringHelper {}"
+        )
+        myFixture.addFileToProject(
+            "com/test/data/CustomClass.java",
+            "package com.test.data; public class CustomClass {}"
+        )
+        val file = myFixture.configureByText(
+            "Test.java",
+            """
+            package com.test;
+
+            public class Test {
+                public void test() {
+                    StringHelper text = new StringHelper();
+                    CustomClass instance = new CustomClass();
+                    ListClass<String> list = new ListClass<>();
+                }
+            }
+            """.trimIndent()
+        )
+        val original = file.text
+
+        val result = AutoImportOrchestrator.previewImports(myFixture.editor)
+
+        assertThat(file.text).isEqualTo(original)
+        assertThat(result).isEqualTo(
+            """
+            package com.test;
+
+            import com.test.data.CustomClass;
+            import com.test.lang.StringHelper;
+            import com.test.util.ListClass;
+
+            public class Test {
+                public void test() {
+                    StringHelper text = new StringHelper();
+                    CustomClass instance = new CustomClass();
+                    ListClass<String> list = new ListClass<>();
+                }
+            }
+            """.trimIndent()
+        )
+    }
+
+    fun testPreviewImportsForKotlinFile() {
+        myFixture.addFileToProject(
+            "com/test/util/ListClass.kt",
+            "package com.test.util; class ListClass<E>"
+        )
+        myFixture.addFileToProject(
+            "com/test/lang/PrintHelper.kt",
+            "package com.test.lang; object PrintHelper { fun println(msg: String) {} }"
+        )
+        myFixture.addFileToProject(
+            "com/test/data/CustomClass.kt",
+            "package com.test.data; class CustomClass"
+        )
+        val file = myFixture.configureByText(
+            "Test.kt",
+            """
+            class Test {
+                fun test() {
+                    PrintHelper.println("hello")
+                    val instance = CustomClass()
+                    val list = ListClass<String>()
+                }
+            }
+            """.trimIndent()
+        )
+        val original = file.text
+
+        val result = AutoImportOrchestrator.previewImports(myFixture.editor)
+
+        assertThat(file.text).isEqualTo(original)
+        assertThat(result).isEqualTo(
+            """
+            import com.test.lang.PrintHelper
+            import com.test.data.CustomClass
+            import com.test.util.ListClass
+
+            class Test {
+                fun test() {
+                    PrintHelper.println("hello")
+                    val instance = CustomClass()
+                    val list = ListClass<String>()
+                }
+            }
+            """.trimIndent()
+        )
+    }
+
+    fun testPreviewImportsWithRangeForJava() {
+        myFixture.addFileToProject(
+            "com/a/Aaa.java",
+            "package com.a; public class Aaa {}"
+        )
+        myFixture.addFileToProject(
+            "com/b/Bbb.java",
+            "package com.b; public class Bbb {}"
+        )
+        myFixture.addFileToProject(
+            "com/c/Ccc.java",
+            "package com.c; public class Ccc {}"
+        )
+
+        val file = myFixture.configureByText(
+            "Client.java",
+            """
+            package client;
+
+            public class Client {
+                public void test() {
+                    Aaa a = new Aaa();
+                    Bbb b = new Bbb();
+                    Ccc c = new Ccc();
+                }
+            }
+            """.trimIndent()
+        )
+        val original = file.text
+
+        val start = original.indexOf("Bbb")
+        val end = start + "Bbb".length
+        val result = AutoImportOrchestrator.previewImports(myFixture.editor, TextRange(start, end))
+
+        assertThat(file.text).isEqualTo(original)
+        assertThat(result).isEqualTo(
+            """
+            package client;
+
+            import com.b.Bbb;
+
+            public class Client {
+                public void test() {
+                    Aaa a = new Aaa();
+                    Bbb b = new Bbb();
+                    Ccc c = new Ccc();
+                }
+            }
+            """.trimIndent()
+        )
+    }
+
+    fun testPreviewImportsWithRangeForKotlin() {
+        myFixture.addFileToProject(
+            "com/a/Aaa.kt",
+            "package com.a; class Aaa"
+        )
+        myFixture.addFileToProject(
+            "com/b/Bbb.kt",
+            "package com.b; class Bbb"
+        )
+        myFixture.addFileToProject(
+            "com/c/Ccc.kt",
+            "package com.c; class Ccc"
+        )
+
+        val file = myFixture.configureByText(
+            "Client.kt",
+            """
+            class Client {
+                fun test() {
+                    val a = Aaa()
+                    val b = Bbb()
+                    val c = Ccc()
+                }
+            }
+            """.trimIndent()
+        )
+        val original = file.text
+
+        val start = original.indexOf("Bbb()")
+        val end = start + "Bbb()".length
+        val result = AutoImportOrchestrator.previewImports(myFixture.editor, TextRange(start, end))
+
+        assertThat(file.text).isEqualTo(original)
+        assertThat(result).isEqualTo(
+            """
+            import com.b.Bbb
+
+            class Client {
+                fun test() {
+                    val a = Aaa()
+                    val b = Bbb()
+                    val c = Ccc()
+                }
+            }
+            """.trimIndent()
+        )
+    }
+
+    fun testReturnsSameContentWhenNoUnresolvedJava() {
+        val file = myFixture.configureByText(
+            "NoUnresolved.java",
+            """
+            package client;
+
+            public class NoUnresolved {
+                public void test() {
+                    int x = 1;
+                    java.util.List<String> list = new java.util.ArrayList<>();
+                }
+            }
+            """.trimIndent()
+        )
+        val original = file.text
+
+        val result = AutoImportOrchestrator.previewImports(myFixture.editor)
+
+        assertThat(result).isEqualTo(original)
+    }
+
+    fun testReturnsSameContentWhenNoUnresolvedKotlin() {
+        val file = myFixture.configureByText(
+            "NoUnresolved.kt",
+            """
+            class NoUnresolved {
+                fun test() {
+                    val x = 1
+                }
+            }
+            """.trimIndent()
+        )
+        val original = file.text
+
+        val result = AutoImportOrchestrator.previewImports(myFixture.editor)
+
+        assertThat(result).isEqualTo(original)
+    }
+
+    fun testReturnsNullForUnsupportedFileType() {
+        val file = myFixture.configureByText(
+            "note.txt",
+            "Just some content with Foo and Bar"
+        )
+        val original = file.text
+
+        val result = AutoImportOrchestrator.previewImports(myFixture.editor)
+
+        assertThat(file.text).isEqualTo(original)
+        assertThat(result).isNull()
+    }
+
+    fun testJavaAvoidsImportForSamePackage() {
+        myFixture.addFileToProject(
+            "com/test/util/SamePkg.java",
+            "package com.test.util; public class SamePkg {}"
+        )
+        myFixture.addFileToProject(
+            "com/other/OtherOne.java",
+            "package com.other; public class OtherOne {}"
+        )
+        val file = myFixture.configureByText(
+            "Test.java",
+            """
+            package com.test.util;
+
+            public class Test {
+                public void test() {
+                    SamePkg a = new SamePkg();
+                    OtherOne o = new OtherOne();
+                }
+            }
+            """.trimIndent()
+        )
+        val original = file.text
+        val result = AutoImportOrchestrator.previewImports(myFixture.editor)
+
+        assertThat(file.text).isEqualTo(original)
+        assertThat(result).isEqualTo(
+            """
+            package com.test.util;
+
+            import com.other.OtherOne;
+
+            public class Test {
+                public void test() {
+                    SamePkg a = new SamePkg();
+                    OtherOne o = new OtherOne();
+                }
+            }
+            """.trimIndent()
+        )
+    }
+
+    fun testDoesNotDuplicateExistingImportsJava() {
+        myFixture.addFileToProject(
+            "com/test/lang/StringHelper.java",
+            "package com.test.lang; public class StringHelper {}"
+        )
+        myFixture.addFileToProject(
+            "com/test/data/CustomClass.java",
+            "package com.test.data; public class CustomClass {}"
+        )
+        val file = myFixture.configureByText(
+            "Test.java",
+            """
+            package com.test;
+
+            import com.test.lang.StringHelper;
+
+            public class Test {
+                public void test() {
+                    StringHelper text = new StringHelper();
+                    CustomClass instance = new CustomClass();
+                }
+            }
+            """.trimIndent()
+        )
+        val original = file.text
+
+        val result = AutoImportOrchestrator.previewImports(myFixture.editor)
+
+        assertThat(file.text).isEqualTo(original)
+        assertThat(result).isEqualTo(
+            """
+            package com.test;
+
+            import com.test.data.CustomClass;
+            import com.test.lang.StringHelper;
+
+            public class Test {
+                public void test() {
+                    StringHelper text = new StringHelper();
+                    CustomClass instance = new CustomClass();
+                }
+            }
+            """.trimIndent()
+        )
+    }
+
+    fun testDoesNotDuplicateExistingImportsKotlin() {
+        myFixture.addFileToProject(
+            "com/test/lang/PrintHelper.kt",
+            "package com.test.lang; object PrintHelper { fun println(msg: String) {} }"
+        )
+        myFixture.addFileToProject(
+            "com/test/data/CustomClass.kt",
+            "package com.test.data; class CustomClass"
+        )
+        val file = myFixture.configureByText(
+            "Test.kt",
+            """
+            import com.test.lang.PrintHelper
+
+            class Test {
+                fun test() {
+                    PrintHelper.println("hello")
+                    val instance = CustomClass()
+                }
+            }
+            """.trimIndent()
+        )
+        val original = file.text
+
+        val result = AutoImportOrchestrator.previewImports(myFixture.editor)
+
+        assertThat(file.text).isEqualTo(original)
+        assertThat(result).isEqualTo(
+            """
+            import com.test.lang.PrintHelper
+            import com.test.data.CustomClass
+
+            class Test {
+                fun test() {
+                    PrintHelper.println("hello")
+                    val instance = CustomClass()
+                }
+            }
+            """.trimIndent()
+        )
+    }
+
+    fun testJavaAmbiguousSimpleNameChoosesAlphabeticalFirst() {
+        myFixture.addFileToProject(
+            "com/a/Dupe.java",
+            "package com.a; public class Dupe {}"
+        )
+        myFixture.addFileToProject(
+            "com/b/Dupe.java",
+            "package com.b; public class Dupe {}"
+        )
+        val file = myFixture.configureByText(
+            "Test.java",
+            """
+            package com.test;
+
+            public class Test {
+                public void test() {
+                    Dupe d = null;
+                }
+            }
+            """.trimIndent()
+        )
+        val original = file.text
+
+        val result = AutoImportOrchestrator.previewImports(myFixture.editor)
+
+        assertThat(file.text).isEqualTo(original)
+        assertThat(result).isEqualTo(
+            """
+            package com.test;
+
+            import com.a.Dupe;
+
+            public class Test {
+                public void test() {
+                    Dupe d = null;
+                }
+            }
+            """.trimIndent()
+        )
+    }
+
+    fun testKotlinAmbiguousSimpleNameChoosesAlphabeticalFirst() {
+        myFixture.addFileToProject(
+            "com/a/Dupe.kt",
+            "package com.a; class Dupe"
+        )
+        myFixture.addFileToProject(
+            "com/b/Dupe.kt",
+            "package com.b; class Dupe"
+        )
+        val file = myFixture.configureByText(
+            "Test.kt",
+            """
+            class Test {
+                fun test() {
+                    val d = Dupe()
+                }
+            }
+            """.trimIndent()
+        )
+        val original = file.text
+
+        val result = AutoImportOrchestrator.previewImports(myFixture.editor)
+
+        assertThat(file.text).isEqualTo(original)
+        assertThat(result).isEqualTo(
+            """
+            import com.a.Dupe
+
+            class Test {
+                fun test() {
+                    val d = Dupe()
+                }
+            }
+            """.trimIndent()
+        )
+    }
+
+    fun testMixedKnownAndUnknownSymbolsJava() {
+        myFixture.addFileToProject(
+            "com/a/Known.java",
+            "package com.a; public class Known {}"
+        )
+        val file = myFixture.configureByText(
+            "Test.java",
+            """
+            package com.test;
+
+            public class Test {
+                public void test() {
+                    Known k = new Known();
+                    UnknownFoo u = null;
+                }
+            }
+            """.trimIndent()
+        )
+        val original = file.text
+
+        val result = AutoImportOrchestrator.previewImports(myFixture.editor)
+
+        assertThat(file.text).isEqualTo(original)
+        assertThat(result).isEqualTo(
+            """
+            package com.test;
+
+            import com.a.Known;
+
+            public class Test {
+                public void test() {
+                    Known k = new Known();
+                    UnknownFoo u = null;
+                }
+            }
+            """.trimIndent()
+        )
+    }
+
+    fun testRangeOutsideAnyReference() {
+        myFixture.addFileToProject(
+            "com/a/Aaa.java",
+            "package com.a; public class Aaa {}"
+        )
+        val file = myFixture.configureByText(
+            "Client.java",
+            """
+            package client;
+
+            public class Client {
+                public void test() {
+                    Aaa a = new Aaa();
+                }
+            }
+            """.trimIndent()
+        )
+        val original = file.text
+
+        val result = AutoImportOrchestrator.previewImports(myFixture.editor, TextRange(0, 1))
+
+        assertThat(result).isEqualTo(original)
+    }
+}