From ae7f5d17db1090ad54e73b97e420ef839a7b7828 Mon Sep 17 00:00:00 2001
From: Carl-Robert <carlrobertoh@gmail.com>
Date: Mon, 27 Nov 2023 01:24:02 +0200
Subject: [PATCH] 262 - Support auto code formatting (#292)

---
 .../configuration/ConfigurationComponent.java | 13 ++++++++
 .../ConfigurationConfigurable.java            |  6 +++-
 .../configuration/ConfigurationState.java     |  9 ++++++
 .../carlrobert/codegpt/util/EditorUtil.java   | 31 ++++++++++++++++---
 .../resources/messages/codegpt.properties     |  5 +--
 5 files changed, 56 insertions(+), 8 deletions(-)

diff --git a/src/main/java/ee/carlrobert/codegpt/settings/configuration/ConfigurationComponent.java b/src/main/java/ee/carlrobert/codegpt/settings/configuration/ConfigurationComponent.java
index 23bbe2e1..4816f209 100644
--- a/src/main/java/ee/carlrobert/codegpt/settings/configuration/ConfigurationComponent.java
+++ b/src/main/java/ee/carlrobert/codegpt/settings/configuration/ConfigurationComponent.java
@@ -43,6 +43,7 @@ public class ConfigurationComponent {
   private final JBTable table;
   private final JBCheckBox openNewTabCheckBox;
   private final JBCheckBox methodNameGenerationCheckBox;
+  private final JBCheckBox autoFormattingCheckBox;
   private final JTextArea systemPromptTextArea;
   private final IntegerField maxTokensField;
   private final JBTextField temperatureField;
@@ -104,12 +105,16 @@ public class ConfigurationComponent {
     methodNameGenerationCheckBox = new JBCheckBox(
         CodeGPTBundle.get("configurationConfigurable.enableMethodNameGeneration.label"),
         configuration.isMethodRefactoringEnabled());
+    autoFormattingCheckBox = new JBCheckBox(
+        CodeGPTBundle.get("configurationConfigurable.autoFormatting.label"),
+        configuration.isAutoFormattingEnabled());
 
     mainPanel = FormBuilder.createFormBuilder()
         .addComponent(tablePanel)
         .addVerticalGap(4)
         .addComponent(openNewTabCheckBox)
         .addComponent(methodNameGenerationCheckBox)
+        .addComponent(autoFormattingCheckBox)
         .addVerticalGap(4)
         .addComponent(new TitledSeparator(
             CodeGPTBundle.get("configurationConfigurable.section.assistant.title")))
@@ -267,6 +272,14 @@ public class ConfigurationComponent {
     methodNameGenerationCheckBox.setSelected(disableMethodNameGeneration);
   }
 
+  public boolean isAutoFormattingEnabled() {
+    return autoFormattingCheckBox.isSelected();
+  }
+
+  public void setAutoFormattingEnabled(boolean enabled) {
+    autoFormattingCheckBox.setSelected(enabled);
+  }
+
   class RevertToDefaultsActionButton extends AnActionButton {
 
     RevertToDefaultsActionButton() {
diff --git a/src/main/java/ee/carlrobert/codegpt/settings/configuration/ConfigurationConfigurable.java b/src/main/java/ee/carlrobert/codegpt/settings/configuration/ConfigurationConfigurable.java
index 301f2d9b..364c66ee 100644
--- a/src/main/java/ee/carlrobert/codegpt/settings/configuration/ConfigurationConfigurable.java
+++ b/src/main/java/ee/carlrobert/codegpt/settings/configuration/ConfigurationConfigurable.java
@@ -40,7 +40,9 @@ public class ConfigurationConfigurable implements Configurable {
         || configurationComponent.isCreateNewChatOnEachAction()
         != configuration.isCreateNewChatOnEachAction()
         || configurationComponent.isMethodNameGenerationEnabled()
-        != configuration.isMethodRefactoringEnabled();
+        != configuration.isMethodRefactoringEnabled()
+        || configurationComponent.isAutoFormattingEnabled()
+        != configuration.isAutoFormattingEnabled();
   }
 
   @Override
@@ -54,6 +56,7 @@ public class ConfigurationConfigurable implements Configurable {
         configurationComponent.isCreateNewChatOnEachAction());
     configuration.setMethodNameGenerationEnabled(
         configurationComponent.isMethodNameGenerationEnabled());
+    configuration.setAutoFormattingEnabled(configurationComponent.isAutoFormattingEnabled());
     EditorActionsUtil.refreshActions();
   }
 
@@ -68,6 +71,7 @@ public class ConfigurationConfigurable implements Configurable {
         configuration.isCreateNewChatOnEachAction());
     configurationComponent.setDisableMethodNameGeneration(
         configuration.isMethodRefactoringEnabled());
+    configurationComponent.setAutoFormattingEnabled(configuration.isAutoFormattingEnabled());
     EditorActionsUtil.refreshActions();
   }
 
diff --git a/src/main/java/ee/carlrobert/codegpt/settings/configuration/ConfigurationState.java b/src/main/java/ee/carlrobert/codegpt/settings/configuration/ConfigurationState.java
index 02cfb369..b3d2d190 100644
--- a/src/main/java/ee/carlrobert/codegpt/settings/configuration/ConfigurationState.java
+++ b/src/main/java/ee/carlrobert/codegpt/settings/configuration/ConfigurationState.java
@@ -23,6 +23,7 @@ public class ConfigurationState implements PersistentStateComponent<Configuratio
   private boolean createNewChatOnEachAction;
   private boolean ignoreGitCommitTokenLimit;
   private boolean methodNameGenerationEnabled = true;
+  private boolean autoFormattingEnabled = true;
   private Map<String, String> tableData = EditorActionsUtil.DEFAULT_ACTIONS;
 
   public static ConfigurationState getInstance() {
@@ -95,4 +96,12 @@ public class ConfigurationState implements PersistentStateComponent<Configuratio
   public void setMethodNameGenerationEnabled(boolean methodNameGenerationEnabled) {
     this.methodNameGenerationEnabled = methodNameGenerationEnabled;
   }
+
+  public boolean isAutoFormattingEnabled() {
+    return autoFormattingEnabled;
+  }
+
+  public void setAutoFormattingEnabled(boolean autoFormattingEnabled) {
+    this.autoFormattingEnabled = autoFormattingEnabled;
+  }
 }
diff --git a/src/main/java/ee/carlrobert/codegpt/util/EditorUtil.java b/src/main/java/ee/carlrobert/codegpt/util/EditorUtil.java
index cd4997c1..a3985d36 100644
--- a/src/main/java/ee/carlrobert/codegpt/util/EditorUtil.java
+++ b/src/main/java/ee/carlrobert/codegpt/util/EditorUtil.java
@@ -14,7 +14,9 @@ import com.intellij.openapi.fileEditor.FileDocumentManager;
 import com.intellij.openapi.fileEditor.FileEditorManager;
 import com.intellij.openapi.project.Project;
 import com.intellij.psi.PsiDocumentManager;
+import com.intellij.psi.codeStyle.CodeStyleManager;
 import com.intellij.testFramework.LightVirtualFile;
+import ee.carlrobert.codegpt.settings.configuration.ConfigurationState;
 import java.time.LocalDateTime;
 import java.time.format.DateTimeFormatter;
 import org.jetbrains.annotations.NotNull;
@@ -88,17 +90,36 @@ public final class EditorUtil {
           var editor = getSelectedEditor(project);
           if (editor != null) {
             var selectionModel = editor.getSelectionModel();
-            editor.getDocument()
-                .replaceString(
-                    selectionModel.getSelectionStart(),
-                    selectionModel.getSelectionEnd(),
-                    text);
+            int startOffset = selectionModel.getSelectionStart();
+            int endOffset = selectionModel.getSelectionEnd();
+            var document = editor.getDocument();
+
+            document.replaceString(startOffset, endOffset, text);
+
+            if (ConfigurationState.getInstance().isAutoFormattingEnabled()) {
+              reformatDocument(project, document, startOffset, endOffset);
+            }
+
             editor.getContentComponent().requestFocus();
             selectionModel.removeSelection();
           }
         })));
   }
 
+  private static void reformatDocument(
+      @NotNull Project project,
+      @NotNull Document document,
+      int startOffset,
+      int endOffset) {
+    var psiDocumentManager = PsiDocumentManager.getInstance(project);
+    psiDocumentManager.commitDocument(document);
+    var psiFile = psiDocumentManager.getPsiFile(document);
+    if (psiFile != null) {
+      CodeStyleManager.getInstance(project)
+          .reformatText(psiFile, startOffset, endOffset);
+    }
+  }
+
   public static void disableHighlighting(@NotNull Project project, Document document) {
     var psiFile = PsiDocumentManager.getInstance(project).getPsiFile(document);
     if (psiFile != null) {
diff --git a/src/main/resources/messages/codegpt.properties b/src/main/resources/messages/codegpt.properties
index f606c3e1..ff06f19b 100644
--- a/src/main/resources/messages/codegpt.properties
+++ b/src/main/resources/messages/codegpt.properties
@@ -66,8 +66,9 @@ configurationConfigurable.table.header.actionColumnLabel=Action
 configurationConfigurable.table.header.promptColumnLabel=Prompt
 configurationConfigurable.table.action.revertToDefaults.text=Revert to Defaults
 configurationConfigurable.table.action.addKeymap.text=Add Shortcut
-configurationConfigurable.openNewTabCheckBox.label=Open new tab on each action
-configurationConfigurable.enableMethodNameGeneration.label=Enable automatic method name generation
+configurationConfigurable.openNewTabCheckBox.label=Open a new chat on each action
+configurationConfigurable.enableMethodNameGeneration.label=Enable method name lookup suggestions
+configurationConfigurable.autoFormatting.label=Enable automatic code formatting
 configurationConfigurable.section.assistant.title=Assistant Configuration
 configurationConfigurable.section.assistant.systemPromptField.label=System prompt:
 configurationConfigurable.section.assistant.systemPromptField.comment=The system message helps to set the behaviour of the assistant