From 40d9fc10792708955e1399380fcbbcec0c9d5305 Mon Sep 17 00:00:00 2001
From: DominikDoom <rehdominik@gmail.com>
Date: Sat, 22 Jul 2023 13:07:06 +0200
Subject: [PATCH] Add ability to undo keyword insertion directly after

---
 javascript/__globals.js       |  6 +++++
 javascript/tagAutocomplete.js | 47 ++++++++++++++++++++++++++++++-----
 2 files changed, 47 insertions(+), 6 deletions(-)

diff --git a/javascript/__globals.js b/javascript/__globals.js
index a066ce2..8be8614 100644
--- a/javascript/__globals.js
+++ b/javascript/__globals.js
@@ -36,6 +36,12 @@ let hideBlocked = false;
 var selectedTag = null;
 var oldSelectedTag = null;
 
+// Lora keyword undo/redo history
+var textBeforeKeywordInsertion = "";
+var textAfterKeywordInsertion = "";
+var lastEditWasKeywordInsertion = false;
+var keywordInsertionUndone = false;
+
 // UMI
 var umiPreviousTags = [];
 
diff --git a/javascript/tagAutocomplete.js b/javascript/tagAutocomplete.js
index 14cdedb..879beb1 100644
--- a/javascript/tagAutocomplete.js
+++ b/javascript/tagAutocomplete.js
@@ -451,8 +451,6 @@ async function insertTextAtCursor(textArea, result, tagword, tabCompletedWithout
             let nameDict = modelKeywordDict.get(result.hash);
             let name = result.text + ".safetensors";
 
-            console.log(name, nameDict);
-
             if (nameDict) {
                 if (nameDict.size > 1)
                     keywords = nameDict.get(name);
@@ -460,10 +458,15 @@ async function insertTextAtCursor(textArea, result, tagword, tabCompletedWithout
                     keywords = nameDict.get("none");
             }
 
-            console.log(keywords);
-
             if (keywords && keywords.length > 0) {
-                newPrompt = `${keywords}, ${newPrompt}`;
+                textBeforeKeywordInsertion = newPrompt;
+                
+                newPrompt = `${keywords}, ${newPrompt}`; // Insert keywords
+                
+                textAfterKeywordInsertion = newPrompt;
+                keywordInsertionUndone = false;
+                setTimeout(() => lastEditWasKeywordInsertion = true, 200)
+                
                 keywordsLength = keywords.length + 2; // +2 for the comma and space
             }
         }
@@ -800,6 +803,37 @@ function rubyTagClicked(node, textBefore, prompt, textArea) {
     textArea.setSelectionRange(startPos, endPos);
 }
 
+// Check if the last edit was the keyword insertion, and catch undo/redo in that case
+function checkKeywordInsertionUndo(textArea, event) {
+    if (!TAC_CFG.modelKeywordCompletion) return;
+
+    switch (event.inputType) {
+        case "historyUndo":
+            if (lastEditWasKeywordInsertion && !keywordInsertionUndone) {
+                keywordInsertionUndone = true;
+                textArea.value = textBeforeKeywordInsertion;
+                updateInput(textArea);
+            }
+            break;
+        case "historyRedo":
+            if (lastEditWasKeywordInsertion && keywordInsertionUndone) {
+                keywordInsertionUndone = false;
+                textArea.value = textAfterKeywordInsertion;
+                updateInput(textArea);
+            }
+        case undefined:
+            // undefined is caused by the updateInput event firing, so we just ignore it
+            break;
+        default:
+            // Everything else deactivates the keyword undo and returns to normal undo behavior
+            lastEditWasKeywordInsertion = false;
+            keywordInsertionUndone = false;
+            textBeforeKeywordInsertion = "";
+            textAfterKeywordInsertion = "";
+            break;
+    }
+}
+
 async function autocomplete(textArea, prompt, fixedTag = null) {
     // Return if the function is deactivated in the UI
     if (!isEnabled()) return;
@@ -1090,9 +1124,10 @@ function addAutocompleteToArea(area) {
         hideResults(area);
 
         // Add autocomplete event listener
-        area.addEventListener('input', () => {
+        area.addEventListener('input', (e) => {
             debounce(autocomplete(area, area.value), TAC_CFG.delayTime);
             updateRuby(area, area.value);
+            checkKeywordInsertionUndo(area, e);
         });
         // Add focusout event listener
         area.addEventListener('focusout', debounce(() => {