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"
	"git.mlow.ca/mlow/lmcli/pkg/tui/shared"
	"git.mlow.ca/mlow/lmcli/pkg/tui/views/chat"
	"git.mlow.ca/mlow/lmcli/pkg/tui/views/conversations"
	tea "github.com/charmbracelet/bubbletea"
)

// Application model
type Model struct {
	shared.State

	state         shared.View
	chat          chat.Model
	conversations conversations.Model
}

func initialModel(ctx *lmcli.Context, values shared.Values) Model {
	m := Model{
		State: shared.State{
			Ctx: ctx,
			Values: &values,
		},
	}

	m.state = shared.StateChat
	m.chat = chat.Chat(m.State)
	m.conversations = conversations.Conversations(m.State)
	return m
}

func (m Model) Init() tea.Cmd {
	return tea.Batch(
		m.conversations.Init(),
		m.chat.Init(),
		func() tea.Msg {
			return shared.MsgViewChange(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 shared.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 shared.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 shared.MsgViewChange:
		m.state = shared.View(msg)
		switch m.state {
		case shared.StateChat:
			m.chat.HandleResize(m.Width, m.Height)
		case shared.StateConversations:
			m.conversations.HandleResize(m.Width, m.Height)
		}
		return m, func() tea.Msg { return shared.MsgViewEnter(struct{}{}) }
	case tea.WindowSizeMsg:
		m.Width, m.Height = msg.Width, msg.Height
	}

	var cmd tea.Cmd
	switch m.state {
	case shared.StateConversations:
		m.conversations, cmd = m.conversations.Update(msg)
	case shared.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 {
	switch m.state {
	case shared.StateConversations:
		return m.conversations.View()
	case shared.StateChat:
		return m.chat.View()
	}
	return ""
}

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