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/api"
	"git.mlow.ca/mlow/lmcli/pkg/lmcli"
	"git.mlow.ca/mlow/lmcli/pkg/tui/model"
	"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"
)

type LaunchOptions struct {
	InitialConversation *api.Conversation
	InitialView         shared.View
}

type Model struct {
	App           *model.AppModel
	view          shared.View
	chat          chat.Model
	conversations conversations.Model
	Width         int
	Height        int
}

func initialModel(ctx *lmcli.Context, opts LaunchOptions) Model {
	m := Model{
		App: &model.AppModel{
			Ctx:          ctx,
			Conversation: opts.InitialConversation,
		},
		view: opts.InitialView,
	}

	sharedData := shared.Shared{}

	m.chat = chat.Chat(m.App, sharedData)
	m.conversations = conversations.Conversations(m.App, sharedData)
	return m
}

func (m Model) Init() tea.Cmd {
	return tea.Batch(
		m.conversations.Init(),
		m.chat.Init(),
		func() tea.Msg {
			return shared.MsgViewChange(m.view)
		},
	)
}

func (m *Model) handleGlobalInput(msg tea.KeyMsg) (bool, tea.Cmd) {
	var cmds []tea.Cmd
	switch m.view {
	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.view = shared.View(msg)
		switch m.view {
		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.view {
	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.view {
	case shared.StateConversations:
		return m.conversations.View()
	case shared.StateChat:
		return m.chat.View()
	}
	return ""
}

type LaunchOption func(*LaunchOptions)

func WithInitialConversation(conv *api.Conversation) LaunchOption {
	return func(opts *LaunchOptions) {
		opts.InitialConversation = conv
	}
}

func WithInitialView(view shared.View) LaunchOption {
	return func(opts *LaunchOptions) {
		opts.InitialView = view
	}
}

func Launch(ctx *lmcli.Context, options ...LaunchOption) error {
	opts := &LaunchOptions{
		InitialView: shared.StateChat,
	}
	for _, opt := range options {
		opt(opts)
	}

	program := tea.NewProgram(initialModel(ctx, *opts), tea.WithAltScreen())
	if _, err := program.Run(); err != nil {
		return fmt.Errorf("Error running program: %v", err)
	}
	return nil
}