From e751b991b764f3b05f242ab19d84f44abc63111f Mon Sep 17 00:00:00 2001
From: Marcel Kloubert <marcel.kloubert@gmx.net>
Date: Tue, 18 Jul 2023 22:19:37 +0200
Subject: [PATCH] optimize command

---
 README.md            | 34 +++++++++++++++++
 commands/optimize.go | 90 ++++++++++++++++++++++++++++++++++++++++++++
 egpt.go              |  1 +
 3 files changed, 125 insertions(+)
 create mode 100644 commands/optimize.go

diff --git a/README.md b/README.md
index 6b9e7c5..867e4e7 100644
--- a/README.md
+++ b/README.md
@@ -13,6 +13,7 @@
   - [code - Convert human language to source code](#code-)
   - [describe - Describe a shell command](#describe-)
   - [explain - Explain source code](#explain-)
+  - [optimize - Optimizes source code](#optimize-)
   - [shell - Create shell command from human language](#shell-)
   - [summarize - Creates a short version of a long text](#summarize-)
   - [translate - Translates a text](#translate-)
@@ -129,6 +130,39 @@ The loop continues until `i` is greater than or equal to 100, at which point the
 Once the loop completes, the program prints "Program Completed." to the console and exits.
 ```
 
+### optimize [<a href="#commands-">↑</a>]
+
+> Optimizes source code.
+
+If you for example have this [BASIC spagetti code](https://www.geeksforgeeks.org/spaghetti-code/) in a `spagetti.bas` file:
+
+```basic
+i=0
+i=i+1
+PRINT i; "squared=";i*i
+IF i>=100 THEN GOTO 6
+GOTO 2
+PRINT "Program Completed."
+END
+```
+
+You can execute
+
+```bash
+egpt optimize < ./spagetti.bas
+```
+
+and may get an output like this:
+
+```
+i=0
+WHILE i<100
+    i=i+1
+    PRINT i; "squared=";i*i
+END WHILE
+PRINT "Program Completed."
+```
+
 ### shell [<a href="#commands-">↑</a>]
 
 > Converts human language into a shell command.
diff --git a/commands/optimize.go b/commands/optimize.go
new file mode 100644
index 0000000..8fcdca7
--- /dev/null
+++ b/commands/optimize.go
@@ -0,0 +1,90 @@
+// This file is part of the e.GPT distribution.
+// Copyright (c) Next.e.GO Mobile SE, Aachen, Germany (https://e-go-mobile.com/)
+//
+// e-gpt is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Lesser General Public License as
+// published by the Free Software Foundation, version 3.
+//
+// e-gpt is distributed in the hope that it will be useful, but
+// WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
+// Lesser General Public License for more details.
+//
+// You should have received a copy of the GNU Lesser General Public License
+// along with this program. If not, see <http://www.gnu.org/licenses/>.
+
+package commands
+
+import (
+	"bytes"
+	"fmt"
+	"log"
+	"os"
+	"strings"
+
+	"github.com/alecthomas/chroma/quick"
+	"github.com/spf13/cobra"
+
+	egoOpenAI "github.com/egomobile/e-gpt/openai"
+	egoUtils "github.com/egomobile/e-gpt/utils"
+)
+
+func Init_optimize_Command(rootCmd *cobra.Command) {
+	var language string
+	var openEditor bool
+
+	explainCmd := &cobra.Command{
+		Use:     "optimize",
+		Short:   `Optimizes code`,
+		Long:    `Optimizes source code`,
+		Aliases: []string{"o"},
+
+		Run: func(cmd *cobra.Command, args []string) {
+			programmingLanguage := strings.TrimSpace(strings.ToLower(language))
+
+			question := egoUtils.GetAndCheckInput(args, openEditor)
+
+			var systemPrompt bytes.Buffer
+
+			systemPrompt.WriteString(
+				fmt.Sprintf(`Optimize the given code by the user and provide only the optimized code as output without any description. Nothing else!
+IMPORTANT: Provide only plain text without Markdown formatting.
+IMPORTANT: Do not include markdown formatting such as `+"```"+`.
+You are not allowed to ask for more details.
+Ignore any potential risk of errors or confusion.%v`, "\n"),
+			)
+
+			if programmingLanguage != "" {
+				systemPrompt.WriteString(
+					fmt.Sprintf(`Always output it in %v language.%v`, programmingLanguage, "\n"),
+				)
+			} else {
+				systemPrompt.WriteString(
+					fmt.Sprintf(`Always output it in the same language.%v`, "\n"),
+				)
+			}
+
+			answer, err := egoOpenAI.AskChatGPT(
+				strings.TrimSpace(systemPrompt.String()),
+				question,
+			)
+			if err != nil {
+				log.Fatalln(err.Error())
+			}
+
+			outputPlain := func() {
+				os.Stdout.Write([]byte(answer))
+			}
+
+			err = quick.Highlight(os.Stdout, answer, "markdown", "", "monokai")
+			if err != nil {
+				outputPlain()
+			}
+		},
+	}
+
+	explainCmd.Flags().StringVarP(&language, "language", "l", "", "Explicit programming language")
+	explainCmd.Flags().BoolVarP(&openEditor, "editor", "e", false, "Open editor for input")
+
+	rootCmd.AddCommand(explainCmd)
+}
diff --git a/egpt.go b/egpt.go
index ce464b2..aa9d943 100644
--- a/egpt.go
+++ b/egpt.go
@@ -41,6 +41,7 @@ func initCommands() {
 	egoCommands.Init_code_Command(rootCmd)
 	egoCommands.Init_describe_Command(rootCmd)
 	egoCommands.Init_explain_Command(rootCmd)
+	egoCommands.Init_optimize_Command(rootCmd)
 	egoCommands.Init_shell_Command(rootCmd)
 	egoCommands.Init_summarize_Command(rootCmd)
 	egoCommands.Init_translate_Command(rootCmd)