From ea116f91fcc831c8cc9bdad3ccdb86bc69e08004 Mon Sep 17 00:00:00 2001
From: Keegan George <kgeorge13@gmail.com>
Date: Wed, 20 Dec 2023 10:46:22 -0800
Subject: [PATCH] FIX: Undo issue (#372)

---
 .../modal/thumbnail-suggestions.gjs           | 12 +++---
 .../after-d-editor/ai-helper-context-menu.hbs |  3 +-
 .../after-d-editor/ai-helper-context-menu.js  | 40 ++++++++++---------
 3 files changed, 30 insertions(+), 25 deletions(-)

diff --git a/assets/javascripts/discourse/components/modal/thumbnail-suggestions.gjs b/assets/javascripts/discourse/components/modal/thumbnail-suggestions.gjs
index 889258df..b56a37d8 100644
--- a/assets/javascripts/discourse/components/modal/thumbnail-suggestions.gjs
+++ b/assets/javascripts/discourse/components/modal/thumbnail-suggestions.gjs
@@ -33,13 +33,15 @@ export default class ThumbnailSuggestions extends Component {
 
   @action
   appendSelectedImages() {
-    const composerValue = this.args.composer?.reply || "";
+    const imageMarkdown = "\n\n" + this.selectedImages.join("\n");
 
-    const newValue = composerValue.concat(
-      "\n\n",
-      this.selectedImages.join("\n")
+    const dEditorInput = document.querySelector(".d-editor-input");
+    dEditorInput.setSelectionRange(
+      dEditorInput.value.length,
+      dEditorInput.value.length
     );
-    this.args.composer.set("reply", newValue);
+    dEditorInput.focus();
+    document.execCommand("insertText", false, imageMarkdown);
     this.args.closeModal();
   }
 
diff --git a/assets/javascripts/discourse/connectors/after-d-editor/ai-helper-context-menu.hbs b/assets/javascripts/discourse/connectors/after-d-editor/ai-helper-context-menu.hbs
index bf6843b9..c0494fea 100644
--- a/assets/javascripts/discourse/connectors/after-d-editor/ai-helper-context-menu.hbs
+++ b/assets/javascripts/discourse/connectors/after-d-editor/ai-helper-context-menu.hbs
@@ -96,7 +96,7 @@
   <Modal::DiffModal
     @confirm={{this.confirmChanges}}
     @diff={{this.diff}}
-    @oldValue={{this.selectedText}}
+    @oldValue={{this.initialValue}}
     @newValue={{this.newSelectedText}}
     @revert={{this.undoAIAction}}
     @closeModal={{fn (mut this.showDiffModal) false}}
@@ -105,7 +105,6 @@
 
 {{#if this.showThumbnailModal}}
   <Modal::ThumbnailSuggestions
-    @composer={{@outletArgs.composer}}
     @thumbnails={{this.thumbnailSuggestions}}
     @closeModal={{fn (mut this.showThumbnailModal) false}}
   />
diff --git a/assets/javascripts/discourse/connectors/after-d-editor/ai-helper-context-menu.js b/assets/javascripts/discourse/connectors/after-d-editor/ai-helper-context-menu.js
index 3f1af2cc..34c04dae 100644
--- a/assets/javascripts/discourse/connectors/after-d-editor/ai-helper-context-menu.js
+++ b/assets/javascripts/discourse/connectors/after-d-editor/ai-helper-context-menu.js
@@ -25,8 +25,6 @@ export default class AiHelperContextMenu extends Component {
   @tracked selectedText = "";
   @tracked newSelectedText;
   @tracked loading = false;
-  @tracked oldEditorValue;
-  @tracked newEditorValue;
   @tracked lastUsedOption = null;
   @tracked showDiffModal = false;
   @tracked showThumbnailModal = false;
@@ -34,6 +32,7 @@ export default class AiHelperContextMenu extends Component {
   @tracked popperPlacement = "top-start";
   @tracked previousMenuState = null;
   @tracked customPromptValue = "";
+  @tracked initialValue = "";
 
   CONTEXT_MENU_STATES = {
     triggers: "TRIGGERS",
@@ -128,6 +127,13 @@ export default class AiHelperContextMenu extends Component {
         )
       : "";
 
+    this.selectionRange = canSelect
+      ? {
+          x: document.activeElement.selectionStart,
+          y: document.activeElement.selectionEnd,
+        }
+      : { x: 0, y: 0 };
+
     if (this.selectedText?.length === 0) {
       this.closeContextMenu();
       return;
@@ -151,12 +157,6 @@ export default class AiHelperContextMenu extends Component {
 
   @bind
   onKeyDown(event) {
-    const cmdOrCtrl = event.ctrlKey || event.metaKey;
-
-    if (cmdOrCtrl && event.key === "z" && this.oldEditorValue) {
-      return this.undoAIAction();
-    }
-
     if (event.key === "Escape") {
       return this.closeContextMenu();
     }
@@ -205,22 +205,26 @@ export default class AiHelperContextMenu extends Component {
   }
 
   _updateSuggestedByAI(data) {
-    const composer = this.args.outletArgs.composer;
-    this.oldEditorValue = this._dEditorInput.value;
     this.newSelectedText = data.suggestions[0];
 
-    this.newEditorValue = this.oldEditorValue.replace(
-      this.selectedText,
-      this.newSelectedText
-    );
-
     if (data.diff) {
       this.diff = data.diff;
     }
-    composer.set("reply", this.newEditorValue);
+
+    this._insertAt(
+      this.selectionRange.x,
+      this.selectionRange.y,
+      this.newSelectedText
+    );
     this.menuState = this.CONTEXT_MENU_STATES.review;
   }
 
+  _insertAt(start, end, text) {
+    this._dEditorInput.setSelectionRange(start, end);
+    this._dEditorInput.focus();
+    document.execCommand("insertText", false, text);
+  }
+
   _toggleLoadingState(loading) {
     if (loading) {
       this._dEditorInput.classList.add("loading");
@@ -335,8 +339,7 @@ export default class AiHelperContextMenu extends Component {
 
   @action
   undoAIAction() {
-    const composer = this.args.outletArgs.composer;
-    composer.set("reply", this.oldEditorValue);
+    document.execCommand("undo", false, null);
     // context menu is prevented from closing when in review state
     // so we change to reset state quickly before closing
     this.menuState = this.CONTEXT_MENU_STATES.resets;
@@ -348,6 +351,7 @@ export default class AiHelperContextMenu extends Component {
     this._toggleLoadingState(true);
     this.lastUsedOption = option;
     this.menuState = this.CONTEXT_MENU_STATES.loading;
+    this.initialValue = this.selectedText;
 
     this._activeAIRequest = ajax("/discourse-ai/ai-helper/suggest", {
       method: "POST",