From 3d8d3b61b3b58557ac9ffde094ccfa109158bfd0 Mon Sep 17 00:00:00 2001
From: Matt Low <matt@mlow.ca>
Date: Sat, 16 Mar 2024 05:49:04 +0000
Subject: [PATCH] tui: add ability to select a message

---
 pkg/tui/tui.go | 73 ++++++++++++++++++++++++++++++++++++++++++++------
 1 file changed, 65 insertions(+), 8 deletions(-)

diff --git a/pkg/tui/tui.go b/pkg/tui/tui.go
index 6025324..dd97530 100644
--- a/pkg/tui/tui.go
+++ b/pkg/tui/tui.go
@@ -47,13 +47,15 @@ type model struct {
 	stopSignal      chan interface{}
 	replyChan       chan models.Message
 	replyChunkChan  chan string
-	err             error
 	persistence     bool // whether we will save new messages in the conversation
+	err             error
 
 	// ui state
-	focus          focusState
-	status         string   // a general status message
-	highlightCache []string // a cache of syntax highlighted message content
+	focus           focusState
+	status          string   // a general status message
+	highlightCache  []string // a cache of syntax highlighted message content
+	messageOffsets  []int
+	selectedMessage int
 
 	// ui elements
 	content viewport.Model
@@ -94,8 +96,8 @@ var (
 	headerStyle    = lipgloss.NewStyle().
 			Background(lipgloss.Color("0"))
 	conversationStyle = lipgloss.NewStyle().
-			MarginTop(1).
-			MarginBottom(1)
+				MarginTop(1).
+				MarginBottom(1)
 	footerStyle = lipgloss.NewStyle().
 			BorderTop(true).
 			BorderStyle(lipgloss.NormalBorder())
@@ -377,6 +379,8 @@ func initialModel(ctx *lmcli.Context, convShortname string) model {
 		stopSignal:     make(chan interface{}),
 		replyChan:      make(chan models.Message),
 		replyChunkChan: make(chan string),
+
+		selectedMessage: -1,
 	}
 
 	m.content = viewport.New(0, 0)
@@ -409,11 +413,44 @@ func initialModel(ctx *lmcli.Context, convShortname string) model {
 	return m
 }
 
+// fraction is the fraction of the total screen height into view the offset
+// should be scrolled into view. 0.5 = items will be snapped to middle of
+// view
+func scrollIntoView(vp *viewport.Model, offset int, fraction float32) {
+	currentOffset := vp.YOffset
+	if offset >= currentOffset && offset < currentOffset+vp.Height {
+		return
+	}
+	distance := currentOffset - offset
+	if distance < 0 {
+		// we should scroll down until it just comes into view
+		vp.SetYOffset(currentOffset - (distance + (vp.Height - int(float32(vp.Height)*fraction))) + 1)
+	} else {
+		// we should scroll up
+		vp.SetYOffset(currentOffset - distance - int(float32(vp.Height)*fraction))
+	}
+}
+
 func (m *model) handleMessagesKey(msg tea.KeyMsg) tea.Cmd {
 	switch msg.String() {
 	case "tab":
 		m.focus = focusInput
+		m.updateContent()
 		m.input.Focus()
+	case "ctrl+k":
+		if m.selectedMessage > 0 && len(m.messages) == len(m.messageOffsets) {
+			m.selectedMessage--
+			m.updateContent()
+			offset := m.messageOffsets[m.selectedMessage]
+			scrollIntoView(&m.content, offset, 0.1)
+		}
+	case "ctrl+j":
+		if m.selectedMessage < len(m.messages)-1 && len(m.messages) == len(m.messageOffsets) {
+			m.selectedMessage++
+			m.updateContent()
+			offset := m.messageOffsets[m.selectedMessage]
+			scrollIntoView(&m.content, offset, 0.1)
+		}
 	}
 	return nil
 }
@@ -422,6 +459,10 @@ func (m *model) handleInputKey(msg tea.KeyMsg) tea.Cmd {
 	switch msg.String() {
 	case "esc":
 		m.focus = focusMessages
+		if m.selectedMessage < 0 || m.selectedMessage >= len(m.messages) {
+			m.selectedMessage = len(m.messages) - 1
+		}
+		m.updateContent()
 		m.input.Blur()
 	case "ctrl+s":
 		userInput := strings.TrimSpace(m.input.Value())
@@ -634,7 +675,12 @@ func (m *model) updateContent() {
 func (m *model) conversationView() string {
 	sb := strings.Builder{}
 	msgCnt := len(m.messages)
+
+	m.messageOffsets = make([]int, len(m.messages))
+	lineCnt := conversationStyle.GetMarginTop()
 	for i, message := range m.messages {
+		m.messageOffsets[i] = lineCnt
+
 		icon := "⚙️"
 		friendly := message.Role.FriendlyRole()
 		style := lipgloss.NewStyle().Bold(true).Faint(true)
@@ -650,18 +696,27 @@ func (m *model) conversationView() string {
 			icon = "🔧"
 		}
 
+		// write message heading with space for content
+		user := style.Render(icon + friendly)
+
 		var saved string
 		if message.ID == 0 {
 			saved = lipgloss.NewStyle().Faint(true).Render(" (not saved)")
 		}
 
-		// write message heading with space for content
-		header := fmt.Sprintf(" %s", style.Render(icon+friendly)+saved)
+		var selectedPrefix string
+		if m.focus == focusMessages && i == m.selectedMessage {
+			selectedPrefix = "> "
+		}
+
+		header := lipgloss.NewStyle().PaddingLeft(1).Render(selectedPrefix + user + saved)
 		sb.WriteString(header)
+		lineCnt += lipgloss.Height(header)
 
 		// TODO: special rendering for tool calls/results?
 		if message.Content != "" {
 			sb.WriteString("\n\n")
+			lineCnt += 1
 
 			// write message contents
 			var highlighted string
@@ -672,10 +727,12 @@ func (m *model) conversationView() string {
 			}
 			contents := messageStyle.Width(m.content.Width).Render(highlighted)
 			sb.WriteString(contents)
+			lineCnt += lipgloss.Height(contents)
 		}
 
 		if i < msgCnt-1 {
 			sb.WriteString("\n\n")
+			lineCnt += 1
 		}
 	}
 	return sb.String()