lmcli/pkg/tui/views/chat/input.go

package chat

import (
	"fmt"
	"strings"

	models "git.mlow.ca/mlow/lmcli/pkg/lmcli/model"
	"git.mlow.ca/mlow/lmcli/pkg/tui/shared"
	tuiutil "git.mlow.ca/mlow/lmcli/pkg/tui/util"
	tea "github.com/charmbracelet/bubbletea"
)

type MessageCycleDirection int

const (
	CycleNext MessageCycleDirection = 1
	CyclePrev MessageCycleDirection = -1
)

func (m *Model) HandleInput(msg tea.KeyMsg) (bool, tea.Cmd) {
	switch m.focus {
	case focusInput:
		consumed, cmd := m.handleInputKey(msg)
		if consumed {
			return true, cmd
		}
	case focusMessages:
		consumed, cmd := m.handleMessagesKey(msg)
		if consumed {
			return true, cmd
		}
	}

	switch msg.String() {
	case "esc":
		if m.waitingForReply {
			m.stopSignal <- struct{}{}
			return true, nil
		}
		return true, func() tea.Msg {
			return shared.MsgViewChange(shared.StateConversations)
		}
	case "ctrl+c":
		if m.waitingForReply {
			m.stopSignal <- struct{}{}
			return true, nil
		}
	case "ctrl+p":
		m.persistence = !m.persistence
		return true, nil
	case "ctrl+t":
		m.showToolResults = !m.showToolResults
		m.rebuildMessageCache()
		m.updateContent()
		return true, nil
	case "ctrl+w":
		m.wrap = !m.wrap
		m.rebuildMessageCache()
		m.updateContent()
		return true, nil
	}
	return false, nil
}

// handleMessagesKey handles input when the messages pane is focused
func (m *Model) handleMessagesKey(msg tea.KeyMsg) (bool, tea.Cmd) {
	switch msg.String() {
	case "tab", "enter":
		m.focus = focusInput
		m.updateContent()
		m.input.Focus()
		return true, nil
	case "e":
		message := m.messages[m.selectedMessage]
		cmd := tuiutil.OpenTempfileEditor("message.*.md", message.Content, "# Edit the message below\n")
		m.editorTarget = selectedMessage
		return true, cmd
	case "ctrl+k":
		if m.selectedMessage > 0 && len(m.messages) == len(m.messageOffsets) {
			m.selectedMessage--
			m.updateContent()
			offset := m.messageOffsets[m.selectedMessage]
			tuiutil.ScrollIntoView(&m.content, offset, m.content.Height/2)
		}
		return true, nil
	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]
			tuiutil.ScrollIntoView(&m.content, offset, m.content.Height/2)
		}
		return true, nil
	case "ctrl+h", "ctrl+l":
		dir := CyclePrev
		if msg.String() == "ctrl+l" {
			dir = CycleNext
		}

		var err error
		var selected *models.Message
		if m.selectedMessage == 0 {
			selected, err = m.cycleSelectedRoot(m.conversation, dir)
			if err != nil {
				return true, shared.WrapError(fmt.Errorf("Could not cycle conversation root: %v", err))
			}
		} else if m.selectedMessage > 0 {
			selected, err = m.cycleSelectedReply(&m.messages[m.selectedMessage-1], dir)
			if err != nil {
				return true, shared.WrapError(fmt.Errorf("Could not cycle reply: %v", err))
			}
		}

		if selected == nil {
			return false, nil
		}

		// Retrieve updated view at this point
		newPath, err := m.State.Ctx.Store.PathToLeaf(selected)
		if err != nil {
			m.State.Err = fmt.Errorf("Could not fetch messages: %v", err)
		}

		m.messages = append(m.messages[:m.selectedMessage], newPath...)
		m.rebuildMessageCache()
		m.updateContent()
		return true, nil
	case "ctrl+r":
		// resubmit the conversation with all messages up until and including the selected message
		if m.waitingForReply || len(m.messages) == 0 {
			return true, nil
		}
		m.messages = m.messages[:m.selectedMessage+1]
		m.messageCache = m.messageCache[:m.selectedMessage+1]
		cmd := m.promptLLM()
		m.updateContent()
		m.content.GotoBottom()
		return true, cmd
	}
	return false, nil
}

// handleInputKey handles input when the input textarea is focused
func (m *Model) handleInputKey(msg tea.KeyMsg) (bool, tea.Cmd) {
	switch msg.String() {
	case "esc":
		m.focus = focusMessages
		if len(m.messages) > 0 {
			if m.selectedMessage < 0 || m.selectedMessage >= len(m.messages) {
				m.selectedMessage = len(m.messages) - 1
			}
			offset := m.messageOffsets[m.selectedMessage]
			tuiutil.ScrollIntoView(&m.content, offset, m.content.Height/2)
		}
		m.updateContent()
		m.input.Blur()
		return true, nil
	case "ctrl+s":
		// TODO: call a "handleSend" function with returns a tea.Cmd
		if m.waitingForReply {
			return false, nil
		}

		input := strings.TrimSpace(m.input.Value())
		if input == "" {
			return true, nil
		}

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

		m.addMessage(models.Message{
			Role:    models.MessageRoleUser,
			Content: input,
		})

		m.input.SetValue("")

		if m.persistence {
			err := m.persistConversation()
			if err != nil {
				return true, shared.WrapError(err)
			}
		}

		cmd := m.promptLLM()
		m.updateContent()
		m.content.GotoBottom()
		return true, cmd
	case "ctrl+e":
		cmd := tuiutil.OpenTempfileEditor("message.*.md", m.input.Value(), "# Edit your input below\n")
		m.editorTarget = input
		return true, cmd
	}
	return false, nil
}