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

import (
	"fmt"
	"strings"
	"time"

	"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"
)

type appState int

const (
	stateChat = iota
	stateConversations
	//stateModelSelect // stateOptions?
	//stateHelp
)

// this struct holds the final rendered content of various UI components, and
// gets populated in the application's Update() method. View() simply composes
// these elements into the final view
type views struct {
	header  string
	content string
	error   string
	input   string
	footer  string
}

type (
	// sent when an error occurs
	msgError error
)

type model struct {
	width  int
	height int

	ctx           *lmcli.Context
	convShortname string

	// application state
	state           appState
	conversations   []models.Conversation
	lastReplies     []models.Message
	conversation    *models.Conversation
	messages        []models.Message
	selectedMessage int
	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
	showToolResults bool     // whether tool calls and results are shown
	messageCache    []string // cache of syntax highlighted and wrapped message content
	messageOffsets  []int

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

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,

		views: &views{},
	}

	m.state = stateChat

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

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

	m.input.Focus()
	m.input.FocusedStyle.CursorLine = lipgloss.NewStyle()
	m.input.FocusedStyle.Base = inputFocusedStyle
	m.input.BlurredStyle.Base = inputBlurredStyle
	m.input.ShowLineNumbers = false

	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
}

func (m model) Init() tea.Cmd {
	cmds := []tea.Cmd{
		textarea.Blink,
		m.spinner.Tick,
		m.waitForChunk(),
		m.waitForReply(),
	}
	switch m.state {
	case stateChat:
		if m.convShortname != "" {
			cmds = append(cmds, m.loadConversation(m.convShortname))
		}
	case stateConversations:
		cmds = append(cmds, m.loadConversations())
	}
	return tea.Batch(cmds...)
}

func (m *model) handleGlobalInput(msg tea.KeyMsg) tea.Cmd {
	switch msg.String() {
	case "ctrl+c":
		if m.waitingForReply {
			m.stopSignal <- ""
			return nil
		} else {
			return tea.Quit
		}
	case "q":
		if m.focus != focusInput {
			return tea.Quit
		}
	default:
		switch m.state {
		case stateChat:
			return m.handleChatInput(msg)
		case stateConversations:
			return m.handleConversationsInput(msg)
		}
	}
	return nil
}

func (m model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
	var cmds []tea.Cmd

	switch msg := msg.(type) {
	case tea.KeyMsg:
		cmd := m.handleGlobalInput(msg)
		if cmd != nil {
			return m, cmd
		}
	case tea.WindowSizeMsg:
		m.content.Width = msg.Width
		m.input.SetWidth(msg.Width - m.input.FocusedStyle.Base.GetHorizontalBorderSize())
		m.rebuildMessageCache()
		m.updateContent()
		m.width = msg.Width
		m.height = msg.Height
	case msgError:
		m.err = msg
	}

	switch m.state {
	case stateConversations:
		cmds = append(cmds, m.handleConversationsUpdate(msg)...)
	case stateChat:
		cmds = append(cmds, m.handleChatUpdate(msg)...)
	}

	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)

	if m.views.header != "" {
		sections = append(sections, m.views.header)
	}

	switch m.state {
	case stateConversations:
		sections = append(sections, m.views.content)
		if m.views.error != "" {
			sections = append(sections, m.views.error)
		}
	case stateChat:
		sections = append(sections, m.views.content)
		if m.views.error != "" {
			sections = append(sections, m.views.error)
		}
		sections = append(sections, m.views.input)
	}

	if m.views.footer != "" {
		sections = append(sections, m.views.footer)
	}

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

func (m *model) headerView() string {
	titleStyle := lipgloss.NewStyle().Bold(true)
	var header string
	switch m.state {
	case stateChat:
		var title string
		if m.conversation != nil && m.conversation.Title != "" {
			title = m.conversation.Title
		} else {
			title = "Untitled"
		}
		title = truncateToCellWidth(title, m.width-headerStyle.GetHorizontalPadding(), "...")
		header = titleStyle.Render(title)
	case stateConversations:
		header = titleStyle.Render("Conversations")
	}
	return headerStyle.Width(m.width).Render(header)
}

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) 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 {
		footer = truncateToCellWidth(footer, m.width, "...")
	}
	return footerStyle.Width(m.width).Render(footer)
}

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

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
}