lmcli/pkg/tui/tui.go

package tui

// The terminal UI for lmcli, launched from the `lmcli chat` command
// TODO:
// - conversation list view
// - change model
// - rename conversation
// - set system prompt
// - system prompt library?

import (
	"context"
	"fmt"
	"strings"
	"time"

	cmdutil "git.mlow.ca/mlow/lmcli/pkg/cmd/util"
	"git.mlow.ca/mlow/lmcli/pkg/lmcli"
	models "git.mlow.ca/mlow/lmcli/pkg/lmcli/model"
	"github.com/charmbracelet/bubbles/spinner"
	"github.com/charmbracelet/bubbles/textarea"
	"github.com/charmbracelet/bubbles/viewport"
	tea "github.com/charmbracelet/bubbletea"
	"github.com/charmbracelet/lipgloss"
	"github.com/muesli/reflow/wordwrap"
)

type focusState int

const (
	focusInput focusState = iota
	focusMessages
)

type editorTarget int

const (
	input editorTarget = iota
	selectedMessage
)

type model struct {
	width  int
	height int

	ctx           *lmcli.Context
	convShortname string

	// application state
	conversation    *models.Conversation
	messages        []models.Message
	waitingForReply bool
	editorTarget    editorTarget
	stopSignal      chan interface{}
	replyChan       chan models.Message
	replyChunkChan  chan string
	persistence     bool // whether we will save new messages in the conversation
	err             error

	// ui state
	focus           focusState
	wrap            bool     // whether message content is wrapped to viewport width
	status          string   // a general status message
	highlightCache  []string // a cache of syntax highlighted message content
	messageOffsets  []int
	selectedMessage int

	// ui elements
	content viewport.Model
	input   textarea.Model
	spinner spinner.Model
}

type message struct {
	role    string
	content string
}

// custom tea.Msg types
type (
	// sent on each chunk received from LLM
	msgResponseChunk string
	// sent when response is finished being received
	msgResponseEnd string
	// a special case of msgError that stops the response waiting animation
	msgResponseError error
	// sent on each completed reply
	msgAssistantReply models.Message
	// sent when a conversation is (re)loaded
	msgConversationLoaded *models.Conversation
	// sent when a new conversation title is set
	msgConversationTitleChanged string
	// send when a conversation's messages are laoded
	msgMessagesLoaded []models.Message
	// sent when an error occurs
	msgError error
)

// styles
var (
	userStyle      = lipgloss.NewStyle().Faint(true).Bold(true).Foreground(lipgloss.Color("10"))
	assistantStyle = lipgloss.NewStyle().Faint(true).Bold(true).Foreground(lipgloss.Color("12"))
	messageStyle   = lipgloss.NewStyle().PaddingLeft(2).PaddingRight(2)
	headerStyle    = lipgloss.NewStyle().
			Background(lipgloss.Color("0"))
	conversationStyle = lipgloss.NewStyle().
				MarginTop(1).
				MarginBottom(1)
	footerStyle = lipgloss.NewStyle().
			BorderTop(true).
			BorderStyle(lipgloss.NormalBorder())
)

func (m model) Init() tea.Cmd {
	return tea.Batch(
		textarea.Blink,
		m.spinner.Tick,
		m.loadConversation(m.convShortname),
		m.waitForChunk(),
		m.waitForReply(),
	)
}

func wrapError(err error) tea.Cmd {
	return func() tea.Msg {
		return msgError(err)
	}
}

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":
			if m.waitingForReply {
				m.stopSignal <- ""
			} else {
				return m, tea.Quit
			}
		case "ctrl+p":
			m.persistence = !m.persistence
		case "ctrl+w":
			m.wrap = !m.wrap
			m.updateContent()
		case "q":
			if m.focus != focusInput {
				return m, tea.Quit
			}
		default:
			var inputHandled tea.Cmd
			switch m.focus {
			case focusInput:
				inputHandled = m.handleInputKey(msg)
			case focusMessages:
				inputHandled = m.handleMessagesKey(msg)
			}
			if inputHandled != nil {
				return m, inputHandled
			}
		}
	case tea.WindowSizeMsg:
		m.width = msg.Width
		m.height = msg.Height
		m.content.Width = msg.Width
		m.content.Height = msg.Height - m.getFixedComponentHeight()
		m.input.SetWidth(msg.Width - 1)
		m.updateContent()
	case msgConversationLoaded:
		m.conversation = (*models.Conversation)(msg)
		cmds = append(cmds, m.loadMessages(m.conversation))
	case msgMessagesLoaded:
		m.setMessages(msg)
		m.updateContent()
	case msgResponseChunk:
		chunk := string(msg)
		last := len(m.messages) - 1
		if last >= 0 && m.messages[last].Role == models.MessageRoleAssistant {
			m.setMessageContents(last, m.messages[last].Content+chunk)
		} else {
			m.addMessage(models.Message{
				Role:    models.MessageRoleAssistant,
				Content: chunk,
			})
		}
		m.updateContent()
		cmds = append(cmds, m.waitForChunk()) // wait for the next chunk
	case msgAssistantReply:
		// the last reply that was being worked on is finished
		reply := models.Message(msg)
		last := len(m.messages) - 1
		if last < 0 {
			panic("Unexpected empty messages handling msgReply")
		}
		m.setMessageContents(last, strings.TrimSpace(m.messages[last].Content))
		if m.messages[last].Role == models.MessageRoleAssistant {
			// the last message was an assistant message, so this is a continuation
			if reply.Role == models.MessageRoleToolCall {
				// update last message rrole to tool call
				m.messages[last].Role = models.MessageRoleToolCall
			}
		} else {
			m.addMessage(reply)
		}

		if m.persistence {
			var err error
			if m.conversation.ID == 0 {
				err = m.ctx.Store.SaveConversation(m.conversation)
			}
			if err != nil {
				cmds = append(cmds, wrapError(err))
			} else {
				cmds = append(cmds, m.persistConversation())
			}
		}

		if m.conversation.Title == "" {
			cmds = append(cmds, m.generateConversationTitle())
		}

		m.updateContent()
		cmds = append(cmds, m.waitForReply())
	case msgResponseEnd:
		m.waitingForReply = false
		last := len(m.messages) - 1
		if last < 0 {
			panic("Unexpected empty messages handling msgResponseEnd")
		}
		m.setMessageContents(last, strings.TrimSpace(m.messages[last].Content))
		m.updateContent()
		m.status = "Press ctrl+s to send"
	case msgResponseError:
		m.waitingForReply = false
		m.status = "Press ctrl+s to send"
		m.err = error(msg)
	case msgConversationTitleChanged:
		title := string(msg)
		m.conversation.Title = title
		if m.persistence {
			err := m.ctx.Store.SaveConversation(m.conversation)
			if err != nil {
				cmds = append(cmds, wrapError(err))
			}
		}
	case msgError:
		m.err = error(msg)
	}

	var cmd tea.Cmd
	m.spinner, cmd = m.spinner.Update(msg)
	if cmd != nil {
		cmds = append(cmds, cmd)
	}

	inputCaptured := false
	m.input, cmd = m.input.Update(msg)
	if cmd != nil {
		inputCaptured = true
		cmds = append(cmds, cmd)
	}

	if !inputCaptured {
		m.content, cmd = m.content.Update(msg)
		if cmd != nil {
			cmds = append(cmds, cmd)
		}
	}

	return m, tea.Batch(cmds...)
}

func (m model) View() string {
	if m.width == 0 {
		// this is the case upon initial startup, but it's also a safe bet that
		// we can just skip rendering if the terminal is really 0 width...
		// without this, the m.*View() functions may crash
		return ""
	}

	sections := make([]string, 0, 6)
	sections = append(sections, m.headerView())
	sections = append(sections, m.contentView())
	error := m.errorView()
	if error != "" {
		sections = append(sections, error)
	}
	sections = append(sections, m.inputView())
	sections = append(sections, m.footerView())

	return lipgloss.JoinVertical(
		lipgloss.Left,
		sections...,
	)
}

// returns the total height of "fixed" components, which are those which don't
// change height dependent on window size.
func (m *model) getFixedComponentHeight() int {
	h := 0
	h += m.input.Height()
	h += lipgloss.Height(m.headerView())
	h += lipgloss.Height(m.footerView())
	errorView := m.errorView()
	if errorView != "" {
		h += lipgloss.Height(errorView)
	}
	return h
}

func (m *model) headerView() string {
	titleStyle := lipgloss.NewStyle().
		PaddingLeft(1).
		PaddingRight(1).
		Bold(true)
	var title string
	if m.conversation != nil && m.conversation.Title != "" {
		title = m.conversation.Title
	} else {
		title = "Untitled"
	}
	part := titleStyle.Render(title)

	return headerStyle.Width(m.width).Render(part)
}

func (m *model) contentView() string {
	return m.content.View()
}

func (m *model) errorView() string {
	if m.err == nil {
		return ""
	}
	return lipgloss.NewStyle().
		Width(m.width).
		AlignHorizontal(lipgloss.Center).
		Bold(true).
		Foreground(lipgloss.Color("1")).
		Render(fmt.Sprintf("%s", m.err))
}

func (m *model) inputView() string {
	return m.input.View()
}

func (m *model) footerView() string {
	segmentStyle := lipgloss.NewStyle().PaddingLeft(1).PaddingRight(1).Faint(true)
	segmentSeparator := "|"

	savingStyle := segmentStyle.Copy().Bold(true)
	saving := ""
	if m.persistence {
		saving = savingStyle.Foreground(lipgloss.Color("2")).Render("✅💾")
	} else {
		saving = savingStyle.Foreground(lipgloss.Color("1")).Render("❌💾")
	}

	status := m.status
	if m.waitingForReply {
		status += m.spinner.View()
	}

	leftSegments := []string{
		saving,
		segmentStyle.Render(status),
	}
	rightSegments := []string{
		segmentStyle.Render(fmt.Sprintf("Model: %s", *m.ctx.Config.Defaults.Model)),
	}

	left := strings.Join(leftSegments, segmentSeparator)
	right := strings.Join(rightSegments, segmentSeparator)

	totalWidth := lipgloss.Width(left) + lipgloss.Width(right)
	remaining := m.width - totalWidth

	var padding string
	if remaining > 0 {
		padding = strings.Repeat(" ", remaining)
	}

	footer := left + padding + right
	if remaining < 0 {
		ellipses := "... "
		// this doesn't work very well, due to trying to trim a string with
		// ansii chars already in it
		footer = footer[:(len(footer)+remaining)-len(ellipses)-3] + ellipses
	}
	return footerStyle.Width(m.width).Render(footer)
}

func initialModel(ctx *lmcli.Context, convShortname string) model {
	m := model{
		ctx:           ctx,
		convShortname: convShortname,
		conversation:  &models.Conversation{},
		persistence:   true,

		stopSignal:     make(chan interface{}),
		replyChan:      make(chan models.Message),
		replyChunkChan: make(chan string),

		wrap:            true,
		selectedMessage: -1,
	}

	m.content = viewport.New(0, 0)

	m.input = textarea.New()
	m.input.CharLimit = 0
	m.input.Placeholder = "Enter a message"

	m.input.FocusedStyle.CursorLine = lipgloss.NewStyle()
	m.input.ShowLineNumbers = false
	m.input.SetHeight(4)
	m.input.Focus()

	m.spinner = spinner.New(spinner.WithSpinner(
		spinner.Spinner{
			Frames: []string{
				".  ",
				".. ",
				"...",
				".. ",
				".  ",
				"   ",
			},
			FPS: time.Second / 3,
		},
	))

	m.waitingForReply = false
	m.status = "Press ctrl+s to send"
	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 "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--
			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)
		}
	case "ctrl+r":
		// resubmit the conversation with all messages up until and including
		// the selected message
		if len(m.messages) == 0 {
			return nil
		}
		m.messages = m.messages[:m.selectedMessage+1]
		m.highlightCache = m.highlightCache[:m.selectedMessage+1]
		m.updateContent()
		m.content.GotoBottom()
		return m.promptLLM()
	}
	return nil
}

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())
		if strings.TrimSpace(userInput) == "" {
			return nil
		}

		if len(m.messages) > 0 && m.messages[len(m.messages)-1].Role == models.MessageRoleUser {
			return wrapError(fmt.Errorf("Can't reply to a user message"))
		}

		reply := models.Message{
			Role:    models.MessageRoleUser,
			Content: userInput,
		}

		if m.persistence {
			var err error
			if m.conversation.ID == 0 {
				err = m.ctx.Store.SaveConversation(m.conversation)
			}
			if err != nil {
				return wrapError(err)
			}

			// ensure all messages up to the one we're about to add are
			// persistent
			cmd := m.persistConversation()
			if cmd != nil {
				return cmd
			}
			// persist our new message, returning with any possible errors
			savedReply, err := m.ctx.Store.AddReply(m.conversation, reply)
			if err != nil {
				return wrapError(err)
			}
			reply = *savedReply
		}

		m.input.SetValue("")
		m.addMessage(reply)

		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
	}
	return nil
}

func (m *model) loadConversation(shortname string) tea.Cmd {
	return func() tea.Msg {
		if shortname == "" {
			return nil
		}
		c, err := m.ctx.Store.ConversationByShortName(shortname)
		if err != nil {
			return msgError(fmt.Errorf("Could not lookup conversation: %v", err))
		}
		if c.ID == 0 {
			return msgError(fmt.Errorf("Conversation not found: %s", shortname))
		}
		return msgConversationLoaded(c)
	}
}

func (m *model) loadMessages(c *models.Conversation) tea.Cmd {
	return func() tea.Msg {
		messages, err := m.ctx.Store.Messages(c)
		if err != nil {
			return msgError(fmt.Errorf("Could not load conversation messages: %v\n", err))
		}
		return msgMessagesLoaded(messages)
	}
}

func (m *model) waitForReply() tea.Cmd {
	return func() tea.Msg {
		return msgAssistantReply(<-m.replyChan)
	}
}

func (m *model) waitForChunk() tea.Cmd {
	return func() tea.Msg {
		return msgResponseChunk(<-m.replyChunkChan)
	}
}

func (m *model) generateConversationTitle() tea.Cmd {
	return func() tea.Msg {
		title, err := cmdutil.GenerateTitle(m.ctx, m.conversation)
		if err != nil {
			return msgError(err)
		}
		return msgConversationTitleChanged(title)
	}
}

func (m *model) promptLLM() tea.Cmd {
	m.waitingForReply = true
	m.status = "Press ctrl+c to cancel"

	return func() tea.Msg {
		completionProvider, err := m.ctx.GetCompletionProvider(*m.ctx.Config.Defaults.Model)
		if err != nil {
			return msgError(err)
		}

		requestParams := models.RequestParameters{
			Model:       *m.ctx.Config.Defaults.Model,
			MaxTokens:   *m.ctx.Config.Defaults.MaxTokens,
			Temperature: *m.ctx.Config.Defaults.Temperature,
			ToolBag:     m.ctx.EnabledTools,
		}

		replyHandler := func(msg models.Message) {
			m.replyChan <- msg
		}

		ctx, cancel := context.WithCancel(context.Background())

		canceled := false
		go func() {
			select {
			case <-m.stopSignal:
				canceled = true
				cancel()
			}
		}()

		resp, err := completionProvider.CreateChatCompletionStream(
			ctx, requestParams, m.messages, replyHandler, m.replyChunkChan,
		)

		if err != nil && !canceled {
			return msgResponseError(err)
		}

		return msgResponseEnd(resp)
	}
}

func (m *model) persistConversation() tea.Cmd {
	existingMessages, err := m.ctx.Store.Messages(m.conversation)
	if err != nil {
		return wrapError(fmt.Errorf("Could not retrieve existing conversation messages while trying to save: %v", err))
	}

	existingById := make(map[uint]*models.Message, len(existingMessages))
	for _, msg := range existingMessages {
		existingById[msg.ID] = &msg
	}

	currentById := make(map[uint]*models.Message, len(m.messages))
	for _, msg := range m.messages {
		currentById[msg.ID] = &msg
	}

	for _, msg := range existingMessages {
		_, ok := currentById[msg.ID]
		if !ok {
			err := m.ctx.Store.DeleteMessage(&msg)
			if err != nil {
				return wrapError(fmt.Errorf("Failed to remove messages: %v", err))
			}
		}
	}

	for i, msg := range m.messages {
		if msg.ID > 0 {
			exist, ok := existingById[msg.ID]
			if ok {
				if msg.Content == exist.Content {
					continue
				}
				// update message when contents don't match that of store
				err := m.ctx.Store.UpdateMessage(&msg)
				if err != nil {
					return wrapError(err)
				}
			} else {
				// this would be quite odd... and I'm not sure how to handle
				// it at the time of writing this
			}
		} else {
			newMessage, err := m.ctx.Store.AddReply(m.conversation, msg)
			if err != nil {
				return wrapError(err)
			}
			m.setMessage(i, *newMessage)
		}
	}
	return nil
}

func (m *model) setMessages(messages []models.Message) {
	m.messages = messages
	m.highlightCache = make([]string, len(messages))
	for i, msg := range m.messages {
		highlighted, _ := m.ctx.Chroma.HighlightS(msg.Content)
		m.highlightCache[i] = highlighted
	}
}

func (m *model) setMessage(i int, msg models.Message) {
	if i >= len(m.messages) {
		panic("i out of range")
	}
	highlighted, _ := m.ctx.Chroma.HighlightS(msg.Content)
	m.messages[i] = msg
	m.highlightCache[i] = highlighted
}

func (m *model) addMessage(msg models.Message) {
	highlighted, _ := m.ctx.Chroma.HighlightS(msg.Content)
	m.messages = append(m.messages, msg)
	m.highlightCache = append(m.highlightCache, highlighted)
}

func (m *model) setMessageContents(i int, content string) {
	if i >= len(m.messages) {
		panic("i out of range")
	}
	highlighted, _ := m.ctx.Chroma.HighlightS(content)
	m.messages[i].Content = content
	m.highlightCache[i] = highlighted
}

func (m *model) updateContent() {
	atBottom := m.content.AtBottom()
	m.content.SetContent(m.conversationView())
	if atBottom {
		// if we were at bottom before the update, scroll with the output
		m.content.GotoBottom()
	}
}

// render the conversation into a string
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)

		switch message.Role {
		case models.MessageRoleUser:
			icon = ""
			style = userStyle
		case models.MessageRoleAssistant:
			icon = ""
			style = assistantStyle
		case models.MessageRoleToolCall, models.MessageRoleToolResult:
			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)")
		}

		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
			if m.highlightCache[i] == "" {
				highlighted = message.Content
			} else {
				highlighted = m.highlightCache[i]
			}
			var contents string
			if m.wrap {
				wrapWidth := m.content.Width - messageStyle.GetHorizontalPadding() - 2
				wrapped := wordwrap.String(highlighted, wrapWidth)
				contents = wrapped
			} else {
				contents = highlighted
			}
			sb.WriteString(messageStyle.Width(0).Render(contents))
			lineCnt += lipgloss.Height(contents)
		}

		if i < msgCnt-1 {
			sb.WriteString("\n\n")
			lineCnt += 1
		}
	}
	return conversationStyle.Render(sb.String())
}

func Launch(ctx *lmcli.Context, convShortname string) error {
	p := tea.NewProgram(initialModel(ctx, convShortname), tea.WithAltScreen())
	if _, err := p.Run(); err != nil {
		return fmt.Errorf("Error running program: %v", err)
	}
	return nil
}