feat(code-completion): refactor code completion task and add InsertUtil

#29

Refactored the CodeCompletionTask to improve code readability and maintainability. Added a new utility class, InsertUtil, to handle text insertion and replacement in the document. Also, made minor changes to SelectElementStrategy and LocationInteractionContext to support the refactored code completion task.

core/src/main/kotlin/com/phodal/shirecore/middleware/select/SelectElementStrategy.kt

@@ -1,8 +1,6 @@
 package com.phodal.shirecore.middleware.select
 
 import com.intellij.openapi.application.runInEdt
-import com.intellij.openapi.application.runReadAction
-import com.intellij.openapi.application.runWriteAction
 import com.intellij.openapi.diagnostic.logger
 import com.intellij.openapi.editor.Editor
 import com.intellij.openapi.project.Project
@@ -131,6 +129,12 @@ sealed class SelectElementStrategy {
         }
 
         fun resolvePsiElement(myProject: Project, editor: Editor): PsiElement? {
+            val elementToAction = DefaultPsiElementStrategy().getElementToAction(myProject, editor)
+
+            if (elementToAction != null) {
+                return elementToAction
+            }
+
             val element: PsiElement? = try {
                 editor.caretModel.currentCaret.offset.let {
                     val psiFile = PsiUtilBase.getPsiFileInEditor(editor, myProject) ?: return@let null

core/src/main/kotlin/com/phodal/shirecore/provider/ide/LocationInteractionContext.kt

@@ -2,6 +2,7 @@ package com.phodal.shirecore.provider.ide
 
 import com.intellij.openapi.editor.Editor
 import com.intellij.openapi.project.Project
+import com.intellij.psi.PsiElement
 import com.phodal.shirecore.action.ShireActionLocation
 import com.phodal.shirecore.agent.InteractionType
 import kotlinx.coroutines.flow.Flow
@@ -25,4 +26,6 @@ data class LocationInteractionContext(
      * the [com.phodal.shirecore.llm.ChatMessage]
      */
     val prompt: String,
+
+    val selectElement: PsiElement? = null,
 )

core/src/main/kotlin/com/phodal/shirecore/provider/impl/CodeCompletionTask.kt

@@ -2,10 +2,8 @@ package com.phodal.shirecore.provider.impl
 
 
 import com.intellij.openapi.application.invokeLater
-import com.intellij.openapi.command.WriteCommandAction
 import com.intellij.openapi.diagnostic.logger
 import com.intellij.openapi.editor.Document
-import com.intellij.openapi.editor.ScrollType
 import com.intellij.openapi.progress.ProgressIndicator
 import com.intellij.openapi.progress.Task
 import com.intellij.openapi.project.Project
@@ -15,8 +13,11 @@ import com.intellij.psi.codeStyle.CodeStyleManager
 import com.phodal.shirecore.ShireCoreBundle
 import com.phodal.shirecore.ShireCoroutineScope
 import com.phodal.shirecore.llm.LlmProvider
+import com.phodal.shirecore.markdown.Code
 import com.phodal.shirecore.provider.impl.dto.CodeCompletionRequest
+import kotlinx.coroutines.cancel
 import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.cancellable
 import kotlinx.coroutines.launch
 import kotlin.jvm.internal.Ref
 
@@ -27,9 +28,10 @@ class CodeCompletionTask(private val request: CodeCompletionRequest) :
     private val writeActionGroupId = "code.complete.intention.write.action"
     private val codeMessage = ShireCoreBundle.message("intentions.chat.code.complete.name")
 
-//    private val chunksString = request.element?.let { SimilarChunksWithPaths.createQuery(it, 60) }
+    //    private val chunksString = request.element?.let { SimilarChunksWithPaths.createQuery(it, 60) }
 //    private val commenter = request.element?.let { LanguageCommenters.INSTANCE.forLanguage(it.language) }
 //    private val commentPrefix = commenter?.lineCommentPrefix
+    private var isCanceled: Boolean = false
 
     override fun run(indicator: ProgressIndicator) {
         val prompt = promptText()
@@ -39,27 +41,34 @@ class CodeCompletionTask(private val request: CodeCompletionRequest) :
 
         val editor = request.editor
         ShireCoroutineScope.scope(request.project).launch {
-            val currentOffset = Ref.IntRef()
-            currentOffset.element = request.offset
+            var currentOffset = request.offset
 
             val project = request.project
-            val finalOutput = StringBuilder()
+            val suggestion = StringBuilder()
 
-            flow.collect {
-                finalOutput.append(it)
-                invokeLater {
-                    WriteCommandAction.runWriteCommandAction(project, codeMessage, writeActionGroupId, {
-                        insertStringAndSaveChange(project, it, editor.document, currentOffset.element, false)
-                    })
+            flow.cancellable().collect { char ->
+                if (isCanceled) {
+                    cancel()
+                    return@collect
+                }
+
+                val parsedContent = Code.parse(char).text;
 
-                    currentOffset.element += it.length
-                    editor.caretModel.moveToOffset(currentOffset.element)
-                    editor.scrollingModel.scrollToCaret(ScrollType.MAKE_VISIBLE)
+                suggestion.append(parsedContent)
+                invokeLater {
+                    if (!isCanceled && !request.isReplacement) {
+                        InsertUtil.insertStreamingToDoc(project, parsedContent, editor, currentOffset)
+                        currentOffset += char.length
+                    }
                 }
             }
 
-            logger.info("Suggestion: $finalOutput")
-            request.postExecute?.invoke(finalOutput.toString())
+            if (request.isReplacement) {
+                InsertUtil.replaceText(project, editor, request.element, suggestion.toString())
+            }
+
+            logger.info("Suggestion: $suggestion")
+            request.postExecute?.invoke(suggestion.toString())
         }
     }
 
@@ -86,7 +95,7 @@ class CodeCompletionTask(private val request: CodeCompletionRequest) :
             suggestion: String,
             document: Document,
             startOffset: Int,
-            withReformat: Boolean
+            withReformat: Boolean,
         ) {
             document.insertString(startOffset, suggestion)
             PsiDocumentManager.getInstance(project).commitDocument(document)

core/src/main/kotlin/com/phodal/shirecore/provider/impl/InsertUtil.kt

@@ -0,0 +1,58 @@
+package com.phodal.shirecore.provider.impl
+
+import com.intellij.openapi.application.runReadAction
+import com.intellij.openapi.command.WriteCommandAction
+import com.intellij.openapi.editor.Document
+import com.intellij.openapi.editor.Editor
+import com.intellij.openapi.editor.ScrollType
+import com.intellij.openapi.project.Project
+import com.intellij.openapi.util.TextRange
+import com.intellij.psi.PsiDocumentManager
+import com.intellij.psi.PsiElement
+import com.intellij.psi.codeStyle.CodeStyleManager
+import com.phodal.shirecore.ShireCoreBundle
+
+object InsertUtil {
+    fun insertStringAndSaveChange(
+        project: Project,
+        content: String,
+        document: Document,
+        startOffset: Int,
+        withReformat: Boolean,
+    ) {
+        document.insertString(startOffset, content)
+        PsiDocumentManager.getInstance(project).commitDocument(document)
+
+        if (!withReformat) return
+
+        val psiFile = PsiDocumentManager.getInstance(project).getPsiFile(document)
+        psiFile?.let { file ->
+            val reformatRange = TextRange(startOffset, startOffset + content.length)
+            CodeStyleManager.getInstance(project).reformatText(file, listOf(reformatRange))
+        }
+    }
+
+    fun insertStreamingToDoc(project: Project, char: String, editor: Editor, currentOffset: Int) {
+        WriteCommandAction.runWriteCommandAction(
+            project,
+            ShireCoreBundle.message("intentions.chat.code.complete.name"),
+            "intentions.write.action",
+            {
+                insertStringAndSaveChange(project, char, editor.document, currentOffset, false)
+            })
+
+        editor.caretModel.moveToOffset(currentOffset + char.length)
+        editor.scrollingModel.scrollToCaret(ScrollType.MAKE_VISIBLE)
+    }
+
+    fun replaceText(project: Project, editor: Editor, element: PsiElement?, output: String) {
+        val primaryCaret = editor.caretModel.primaryCaret;
+        val start = runReadAction { primaryCaret.selectionStart }
+        val end = runReadAction { primaryCaret.selectionEnd }
+        val document = runReadAction { editor.document }
+
+        WriteCommandAction.runWriteCommandAction(project) {
+            document.replaceString(start, end, output)
+        }
+    }
+}

core/src/main/resources/messages/ShireCoreBundle.properties

@@ -3,4 +3,4 @@ progress.run.task=Running task
 shire.ref.loading=Loading
 schema.custom-agent.json.display.name=Shire CustomAgent Schema
 intentions.request.background.process.title=Your LLM handle generate file
-intentions.chat.code.complete.name=Code complete
+intentions.chat.code.complete.name=Code Complete