lmcli/pkg/tui/tui.go

package tui

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

import (
	"fmt"

	"git.mlow.ca/mlow/lmcli/pkg/lmcli"
	tea "github.com/charmbracelet/bubbletea"
	"github.com/charmbracelet/lipgloss"
)

type state 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 output
type views struct {
	header  string
	content string
	error   string
	input   string
	footer  string
}

type (
	// send to change the current state
	msgStateChange state
	// sent to a state when it is entered
	msgStateEnter struct{}
	// sent when an error occurs
	msgError error
)

type Options struct {
	convShortname string
}

type basemodel struct {
	opts   *Options
	ctx    *lmcli.Context
	views  *views
	err    error
	width  int
	height int
}

type model struct {
	basemodel

	state         state
	chat          chatModel
	conversations conversationsModel
}

func initialModel(ctx *lmcli.Context, opts Options) model {
	m := model{
		basemodel: basemodel{
			opts:  &opts,
			ctx:   ctx,
			views: &views{},
		},
	}
	m.state = stateChat
	m.chat = newChatModel(&m)
	m.conversations = newConversationsModel(&m)
	return m
}

func (m model) Init() tea.Cmd {
	return tea.Batch(
		m.conversations.Init(),
		m.chat.Init(),
		func() tea.Msg {
			return msgStateChange(m.state)
		},
	)
}

func (m *model) handleGlobalInput(msg tea.KeyMsg) (bool, tea.Cmd) {
	// delegate input to the active child state first, only handling it at the
	// global level if the child state does not
	var cmds []tea.Cmd
	switch m.state {
	case stateChat:
		handled, cmd := m.chat.handleInput(msg)
		cmds = append(cmds, cmd)
		if handled {
			m.chat, cmd = m.chat.Update(nil)
			cmds = append(cmds, cmd)
			return true, tea.Batch(cmds...)
		}
	case stateConversations:
		handled, cmd := m.conversations.handleInput(msg)
		cmds = append(cmds, cmd)
		if handled {
			m.conversations, cmd = m.conversations.Update(nil)
			cmds = append(cmds, cmd)
			return true, tea.Batch(cmds...)
		}
	}
	switch msg.String() {
	case "ctrl+c", "ctrl+q":
		return true, tea.Quit
	}
	return false, nil
}

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

	switch msg := msg.(type) {
	case tea.KeyMsg:
		handled, cmd := m.handleGlobalInput(msg)
		if handled {
			return m, cmd
		}
	case msgStateChange:
		m.state = state(msg)
		switch m.state {
		case stateChat:
			m.chat.handleResize(m.width, m.height)
		case stateConversations:
			m.conversations.handleResize(m.width, m.height)
		}
		return m, func() tea.Msg { return msgStateEnter(struct{}{}) }
	case msgConversationSelected:
		// passed up through conversation list model
		m.opts.convShortname = msg.ShortName.String
		cmds = append(cmds, func() tea.Msg {
			return msgStateChange(stateChat)
		})
	case tea.WindowSizeMsg:
		m.width, m.height = msg.Width, msg.Height
	case msgError:
		m.err = msg
	}

	var cmd tea.Cmd
	switch m.state {
	case stateConversations:
		m.conversations, cmd = m.conversations.Update(msg)
	case stateChat:
		m.chat, cmd = m.chat.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)

	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 errorBanner(err error, width int) string {
	if err == nil {
		return ""
	}
	return lipgloss.NewStyle().
		Width(width).
		AlignHorizontal(lipgloss.Center).
		Bold(true).
		Foreground(lipgloss.Color("1")).
		Render(fmt.Sprintf("%s", err))
}

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, Options{convShortname}), tea.WithAltScreen())
	if _, err := p.Run(); err != nil {
		return fmt.Errorf("Error running program: %v", err)
	}
	return nil
}