From a8ffdc156ae00f97392c0fc232f746e448b8afb8 Mon Sep 17 00:00:00 2001
From: Matt Low <matt@mlow.ca>
Date: Sat, 16 Mar 2024 21:56:45 +0000
Subject: [PATCH] tui: open input/messages for editing in $EDITOR

---
 pkg/tui/tui.go  | 36 ++++++++++++++++++++++++++++++++++--
 pkg/tui/util.go | 42 ++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 76 insertions(+), 2 deletions(-)
 create mode 100644 pkg/tui/util.go

diff --git a/pkg/tui/tui.go b/pkg/tui/tui.go
index dd97530..7f9153c 100644
--- a/pkg/tui/tui.go
+++ b/pkg/tui/tui.go
@@ -2,8 +2,7 @@ package tui
 
 // The terminal UI for lmcli, launched from the `lmcli chat` command
 // TODO:
-// - binding to open selected message/input in $EDITOR
-// - ability to continue or retry previous response
+// - ability to continue an incomplete or missing assistant response
 // - conversation list view
 // - change model
 // - rename conversation
@@ -33,6 +32,13 @@ const (
 	focusMessages
 )
 
+type editorTarget int
+
+const (
+	input editorTarget = iota
+	selectedMessage
+)
+
 type model struct {
 	width  int
 	height int
@@ -44,6 +50,7 @@ type model struct {
 	conversation    *models.Conversation
 	messages        []models.Message
 	waitingForReply bool
+	editorTarget    editorTarget
 	stopSignal      chan interface{}
 	replyChan       chan models.Message
 	replyChunkChan  chan string
@@ -123,6 +130,22 @@ func (m model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
 	var cmds []tea.Cmd
 
 	switch msg := msg.(type) {
+	case msgTempfileEditorClosed:
+		contents := string(msg)
+		switch m.editorTarget {
+		case input:
+			m.input.SetValue(contents)
+		case selectedMessage:
+			m.setMessageContents(m.selectedMessage, contents)
+			if m.persistence && m.messages[m.selectedMessage].ID > 0 {
+				// update persisted message
+				err := m.ctx.Store.UpdateMessage(&m.messages[m.selectedMessage])
+				if err != nil {
+					cmds = append(cmds, wrapError(fmt.Errorf("Could not save edited message: %v", err)))
+				}
+			}
+			m.updateContent()
+		}
 	case tea.KeyMsg:
 		switch msg.String() {
 		case "ctrl+c":
@@ -437,6 +460,11 @@ func (m *model) handleMessagesKey(msg tea.KeyMsg) tea.Cmd {
 		m.focus = focusInput
 		m.updateContent()
 		m.input.Focus()
+	case "e":
+		message := m.messages[m.selectedMessage]
+		cmd := openTempfileEditor("message.*.md", message.Content, "# Edit the message below\n")
+		m.editorTarget = selectedMessage
+		return cmd
 	case "ctrl+k":
 		if m.selectedMessage > 0 && len(m.messages) == len(m.messageOffsets) {
 			m.selectedMessage--
@@ -508,6 +536,10 @@ func (m *model) handleInputKey(msg tea.KeyMsg) tea.Cmd {
 		m.updateContent()
 		m.content.GotoBottom()
 		return m.promptLLM()
+	case "ctrl+e":
+		cmd := openTempfileEditor("message.*.md", m.input.Value(), "# Edit your input below\n")
+		m.editorTarget = input
+		return cmd
 	case "ctrl+r":
 		if len(m.messages) == 0 {
 			return nil
diff --git a/pkg/tui/util.go b/pkg/tui/util.go
new file mode 100644
index 0000000..81eac8f
--- /dev/null
+++ b/pkg/tui/util.go
@@ -0,0 +1,42 @@
+package tui
+
+import (
+	"os"
+	"os/exec"
+	"strings"
+
+	tea "github.com/charmbracelet/bubbletea"
+)
+
+type msgTempfileEditorClosed string
+
+// openTempfileEditor opens an $EDITOR on a new temporary file with the given
+// content. Upon closing, the contents of the file are read back returned
+// wrapped in a msgTempfileEditorClosed returned by the tea.Cmd
+func openTempfileEditor(pattern string, content string, placeholder string) tea.Cmd {
+	msgFile, _ := os.CreateTemp("/tmp", pattern)
+
+	err := os.WriteFile(msgFile.Name(), []byte(placeholder+content), os.ModeAppend)
+	if err != nil {
+		return wrapError(err)
+	}
+
+	editor := os.Getenv("EDITOR")
+	if editor == "" {
+		editor = "vim"
+	}
+
+	c := exec.Command(editor, msgFile.Name())
+	return tea.ExecProcess(c, func(err error) tea.Msg {
+		bytes, err := os.ReadFile(msgFile.Name())
+		if err != nil {
+			return msgError(err)
+		}
+		fileContents := string(bytes)
+		if strings.HasPrefix(fileContents, placeholder) {
+			fileContents = fileContents[len(placeholder):]
+		}
+		stripped := strings.Trim(fileContents, "\n \t")
+		return msgTempfileEditorClosed(stripped)
+	})
+}