From bd0977ce75d5d755b6dce087b8dd5c01ecb4ce91 Mon Sep 17 00:00:00 2001
From: Carl-Robert Linnupuu <carlrobertoh@gmail.com>
Date: Thu, 12 Sep 2024 02:16:12 +0300
Subject: [PATCH] fix: highlighted text referencing in prompts

---
 .../toolwindow/chat/ChatToolWindowTabPanel.java   | 15 +++++++++++++--
 1 file changed, 13 insertions(+), 2 deletions(-)

diff --git a/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/ChatToolWindowTabPanel.java b/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/ChatToolWindowTabPanel.java
index 2a3b5c15..9e7f7ce2 100644
--- a/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/ChatToolWindowTabPanel.java
+++ b/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/ChatToolWindowTabPanel.java
@@ -303,8 +303,19 @@ public class ChatToolWindowTabPanel implements Disposable {
     }
     promptBuilder.append(remainingText);
 
-    message.setUserMessage(promptBuilder.toString());
-    message.setPrompt(promptBuilder.toString());
+    String highlightedTextMd = "";
+    if (editor != null) {
+      var selectionModel = editor.getSelectionModel();
+      var selectedText = selectionModel.getSelectedText();
+      if (selectedText != null && !selectedText.isEmpty()) {
+        var fileExtension = FileUtil.getFileExtension(editor.getVirtualFile().getName());
+        highlightedTextMd = format("\n```%s\n%s\n```\n", fileExtension, selectedText);
+        selectionModel.removeSelection();
+      }
+    }
+
+    message.setUserMessage(highlightedTextMd + promptBuilder);
+    message.setPrompt(highlightedTextMd + promptBuilder);
 
     sendMessage(message, ConversationType.DEFAULT, processEditorSelection(editor, message));
     return Unit.INSTANCE;