From ee4aee2a571dee59bcd3379a2d18ffe42e7566fe Mon Sep 17 00:00:00 2001
From: Carl-Robert Linnupuu <carlrobertoh@gmail.com>
Date: Mon, 25 Nov 2024 11:30:09 +0000
Subject: [PATCH] fix: remaining completion when code completion overtyped

---
 .../CodeCompletionInsertAction.kt             |  7 +++++++
 .../CodeCompletionSuggestionUpdateAdapter.kt  | 20 +++++++++++++++++++
 2 files changed, 27 insertions(+)

diff --git a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionInsertAction.kt b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionInsertAction.kt
index 96d3cdfc..ebc289c4 100644
--- a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionInsertAction.kt
+++ b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionInsertAction.kt
@@ -14,6 +14,7 @@ import com.intellij.openapi.editor.actionSystem.EditorAction
 import com.intellij.openapi.editor.actionSystem.EditorWriteActionHandler
 import com.intellij.psi.PsiDocumentManager
 import com.intellij.util.concurrency.ThreadingAssertions
+import ee.carlrobert.codegpt.CodeGPTKeys
 import ee.carlrobert.codegpt.CodeGPTKeys.REMAINING_EDITOR_COMPLETION
 
 class CodeCompletionInsertAction :
@@ -31,6 +32,12 @@ class CodeCompletionInsertAction :
                 .map { it.element as CodeCompletionTextElement }
 
             if (elements.isEmpty()) {
+                val textToInsert = context.textToInsert()
+                val remainingCompletion = REMAINING_EDITOR_COMPLETION.get(editor)
+                if (remainingCompletion.isNotEmpty()) {
+                    REMAINING_EDITOR_COMPLETION.set(editor, remainingCompletion.removePrefix(textToInsert))
+                }
+
                 InlineCompletion.getHandlerOrNull(editor)?.insert()
                 return
             }
diff --git a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionSuggestionUpdateAdapter.kt b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionSuggestionUpdateAdapter.kt
index d423b254..d48c31b6 100644
--- a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionSuggestionUpdateAdapter.kt
+++ b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionSuggestionUpdateAdapter.kt
@@ -7,10 +7,20 @@ import com.intellij.codeInsight.inline.completion.suggestion.InlineCompletionSug
 import com.intellij.codeInsight.inline.completion.suggestion.InlineCompletionSuggestionUpdateManager.UpdateResult.Changed
 import com.intellij.codeInsight.inline.completion.suggestion.InlineCompletionSuggestionUpdateManager.UpdateResult.Invalidated
 import com.intellij.codeInsight.inline.completion.suggestion.InlineCompletionVariant
+import com.intellij.openapi.editor.Editor
+import ee.carlrobert.codegpt.CodeGPTKeys.REMAINING_EDITOR_COMPLETION
 
 class CodeCompletionSuggestionUpdateAdapter :
     InlineCompletionSuggestionUpdateManager.Default() {
 
+    override fun onDocumentChange(
+        event: InlineCompletionEvent.DocumentChange,
+        variant: InlineCompletionVariant.Snapshot
+    ): UpdateResult {
+        updateRemainingCompletion(event.editor, event.typing.typed)
+        return super.onDocumentChange(event, variant)
+    }
+
     override fun onCustomEvent(
         event: InlineCompletionEvent,
         variant: InlineCompletionVariant.Snapshot
@@ -23,10 +33,20 @@ class CodeCompletionSuggestionUpdateAdapter :
         val textToInsert = event.toRequest().run {
             CompletionSplitter.split(completionText)
         }
+
+        updateRemainingCompletion(event.toRequest().editor, textToInsert)
+
         return Changed(
             variant.copy(
                 listOf(InlineCompletionGrayTextElement(completionText.removePrefix(textToInsert)))
             )
         )
     }
+
+    private fun updateRemainingCompletion(editor: Editor, textToInsert: String) {
+        val remainingCompletion = REMAINING_EDITOR_COMPLETION.get(editor)
+        if (remainingCompletion.isNotEmpty()) {
+            REMAINING_EDITOR_COMPLETION.set(editor, remainingCompletion.removePrefix(textToInsert))
+        }
+    }
 }
\ No newline at end of file