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

package chat

import (
	"context"
	"encoding/json"
	"fmt"
	"strings"
	"time"

	cmdutil "git.mlow.ca/mlow/lmcli/pkg/cmd/util"
	models "git.mlow.ca/mlow/lmcli/pkg/lmcli/model"
	"git.mlow.ca/mlow/lmcli/pkg/tui/shared"
	"git.mlow.ca/mlow/lmcli/pkg/tui/styles"
	tuiutil "git.mlow.ca/mlow/lmcli/pkg/tui/util"
	"github.com/charmbracelet/bubbles/cursor"
	"github.com/charmbracelet/bubbles/spinner"
	"github.com/charmbracelet/bubbles/textarea"
	"github.com/charmbracelet/bubbles/viewport"
	tea "github.com/charmbracelet/bubbletea"
	"github.com/charmbracelet/lipgloss"
	"github.com/muesli/reflow/wordwrap"
	"github.com/muesli/reflow/wrap"
	"gopkg.in/yaml.v2"
)

type focusState int

const (
	focusInput focusState = iota
	focusMessages
)

type editorTarget int

const (
	input editorTarget = iota
	selectedMessage
)

// custom tea.Msg types
type (
	// sent on each chunk received from LLM
	msgResponseChunk string
	// sent when response is finished being received
	msgResponseEnd string
	// a special case of common.MsgError that stops the response waiting animation
	msgResponseError error
	// sent on each completed reply
	msgAssistantReply models.Message
	// sent when a conversation is (re)loaded
	msgConversationLoaded *models.Conversation
	// sent when a new conversation title is set
	msgConversationTitleChanged string
	// sent when a conversation's messages are laoded
	msgMessagesLoaded []models.Message
)

type Model struct {
	shared.State
	width  int
	height int

	// app state
	conversation    *models.Conversation
	rootMessages    []models.Message
	messages        []models.Message
	selectedMessage int
	waitingForReply bool
	editorTarget    editorTarget
	stopSignal      chan struct{}
	replyChan       chan models.Message
	replyChunkChan  chan string
	persistence     bool // whether we will save new messages in the conversation

	tokenCount uint
	startTime  time.Time
	elapsed    time.Duration

	// 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
	replyCursor cursor.Model // cursor to indicate incoming response
}

func Chat(tui shared.State) Model {
	m := Model{
		State: tui,

		conversation: &models.Conversation{},
		persistence:  true,

		stopSignal:     make(chan struct{}),
		replyChan:      make(chan models.Message),
		replyChunkChan: make(chan string),

		wrap:            true,
		selectedMessage: -1,

		content: viewport.New(0, 0),
		input:   textarea.New(),
		spinner: spinner.New(spinner.WithSpinner(
			spinner.Spinner{
				Frames: []string{
					".  ",
					".. ",
					"...",
					".. ",
					".  ",
					"   ",
				},
				FPS: time.Second / 3,
			},
		)),
		replyCursor: cursor.New(),
	}

	m.replyCursor.SetChar(" ")
	m.replyCursor.Focus()

	system := tui.Ctx.GetSystemPrompt()
	if system != "" {
		m.messages = []models.Message{{
			Role:    models.MessageRoleSystem,
			Content: system,
		}}
	}

	m.input.Focus()
	m.input.MaxHeight = 0
	m.input.CharLimit = 0
	m.input.ShowLineNumbers = false
	m.input.Placeholder = "Enter a message"

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

	m.waitingForReply = false
	m.status = "Press ctrl+s to send"
	return m
}

// styles
var (
	messageHeadingStyle = lipgloss.NewStyle().
				MarginTop(1).
				MarginBottom(1).
				PaddingLeft(1).
				Bold(true)

	userStyle = lipgloss.NewStyle().Faint(true).Foreground(lipgloss.Color("10"))

	assistantStyle = lipgloss.NewStyle().Faint(true).Foreground(lipgloss.Color("12"))

	messageStyle = lipgloss.NewStyle().
			PaddingLeft(2).
			PaddingRight(2)

	inputFocusedStyle = lipgloss.NewStyle().
				Border(lipgloss.RoundedBorder(), true, true, true, false)

	inputBlurredStyle = lipgloss.NewStyle().
				Faint(true).
				Border(lipgloss.RoundedBorder(), true, true, true, false)

	footerStyle = lipgloss.NewStyle()
)

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
}

func (m Model) Init() tea.Cmd {
	return tea.Batch(
		m.waitForChunk(),
		m.waitForReply(),
	)
}

func (m *Model) HandleResize(width, height int) {
	m.width, m.height = width, height
	m.content.Width = width
	m.input.SetWidth(width - m.input.FocusedStyle.Base.GetHorizontalFrameSize())
	if len(m.messages) > 0 {
		m.rebuildMessageCache()
		m.updateContent()
	}
}

func (m Model) Update(msg tea.Msg) (Model, tea.Cmd) {
	var cmds []tea.Cmd
	switch msg := msg.(type) {
	case shared.MsgViewEnter:
		// wake up spinners and cursors
		cmds = append(cmds, cursor.Blink, m.spinner.Tick)

		if m.State.Values.ConvShortname != "" && m.conversation.ShortName.String != m.State.Values.ConvShortname {
			cmds = append(cmds, m.loadConversation(m.State.Values.ConvShortname))
		}

		m.rebuildMessageCache()
		m.updateContent()
	case tea.WindowSizeMsg:
		m.HandleResize(msg.Width, msg.Height)
	case tuiutil.MsgTempfileEditorClosed:
		contents := string(msg)
		switch m.editorTarget {
		case input:
			m.input.SetValue(contents)
		case selectedMessage:
			m.setMessageContents(m.selectedMessage, contents)
			if m.persistence && m.messages[m.selectedMessage].ID > 0 {
				// update persisted message
				err := m.State.Ctx.Store.UpdateMessage(&m.messages[m.selectedMessage])
				if err != nil {
					cmds = append(cmds, shared.WrapError(fmt.Errorf("Could not save edited message: %v", err)))
				}
			}
			m.updateContent()
		}
	case msgConversationLoaded:
		m.conversation = (*models.Conversation)(msg)
		m.rootMessages, _ = m.State.Ctx.Store.RootMessages(m.conversation.ID)
		cmds = append(cmds, m.loadMessages(m.conversation))
	case msgMessagesLoaded:
		m.selectedMessage = len(msg) - 1
		m.messages = msg
		m.rebuildMessageCache()
		m.updateContent()
		m.content.GotoBottom()
	case msgResponseChunk:
		chunk := string(msg)
		if chunk == "" {
			break
		}

		last := len(m.messages) - 1
		if last >= 0 && m.messages[last].Role.IsAssistant() {
			// append chunk to existing message
			m.setMessageContents(last, m.messages[last].Content+chunk)
		} else {
			// use chunk in new message
			m.addMessage(models.Message{
				Role:    models.MessageRoleAssistant,
				Content: chunk,
			})
		}
		m.updateContent()
		cmds = append(cmds, m.waitForChunk()) // wait for the next chunk

		// show cursor and reset blink interval (simulate typing)
		m.replyCursor.Blink = false
		cmds = append(cmds, m.replyCursor.BlinkCmd())

		m.tokenCount++
		m.elapsed = time.Now().Sub(m.startTime)
	case msgAssistantReply:
		// the last reply that was being worked on is finished
		reply := models.Message(msg)
		reply.Content = strings.TrimSpace(reply.Content)

		last := len(m.messages) - 1
		if last < 0 {
			panic("Unexpected empty messages handling msgAssistantReply")
		}

		if reply.Role.IsAssistant() && m.messages[last].Role.IsAssistant() {
			// this was a continuation, so replace the previous message with the completed reply
			m.setMessage(last, reply)
		} else {
			m.addMessage(reply)
		}

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

		if m.conversation.Title == "" {
			cmds = append(cmds, m.generateConversationTitle())
		}

		m.updateContent()
		cmds = append(cmds, m.waitForReply())
	case msgResponseEnd:
		m.waitingForReply = false
		last := len(m.messages) - 1
		if last < 0 {
			panic("Unexpected empty messages handling msgResponseEnd")
		}
		m.setMessageContents(last, strings.TrimSpace(m.messages[last].Content))
		m.updateContent()
		m.status = "Press ctrl+s to send"
	case msgResponseError:
		m.waitingForReply = false
		m.status = "Press ctrl+s to send"
		m.State.Err = error(msg)
		m.updateContent()
	case msgConversationTitleChanged:
		title := string(msg)
		m.conversation.Title = title
		if m.persistence {
			err := m.State.Ctx.Store.UpdateConversation(m.conversation)
			if err != nil {
				cmds = append(cmds, shared.WrapError(err))
			}
		}
	case cursor.BlinkMsg:
		if m.waitingForReply {
			// ensure we show updated "wait for response" cursor blink state
			m.updateContent()
		}
	}

	var cmd tea.Cmd
	m.spinner, cmd = m.spinner.Update(msg)
	if cmd != nil {
		cmds = append(cmds, cmd)
	}
	m.replyCursor, cmd = m.replyCursor.Update(msg)
	if cmd != nil {
		cmds = append(cmds, cmd)
	}

	prevInputLineCnt := m.input.LineCount()
	inputCaptured := false
	m.input, cmd = m.input.Update(msg)
	if cmd != nil {
		inputCaptured = true
		cmds = append(cmds, cmd)
	}

	if !inputCaptured {
		m.content, cmd = m.content.Update(msg)
		if cmd != nil {
			cmds = append(cmds, cmd)
		}
	}

	// update views once window dimensions are known
	if m.width > 0 {
		m.State.Views.Header = m.headerView()
		m.State.Views.Footer = m.footerView()
		m.State.Views.Error = tuiutil.ErrorBanner(m.State.Err, m.width)
		fixedHeight := tuiutil.Height(m.State.Views.Header) + tuiutil.Height(m.State.Views.Error) + tuiutil.Height(m.State.Views.Footer)

		// calculate clamped input height to accomodate input text
		// minimum 4 lines, maximum half of content area
		newHeight := max(4, min((m.height-fixedHeight-1)/2, m.input.LineCount()))
		m.input.SetHeight(newHeight)
		m.State.Views.Input = m.input.View()

		// remaining height towards content
		m.content.Height = m.height - fixedHeight - tuiutil.Height(m.State.Views.Input)
		m.State.Views.Content = m.content.View()
	}

	// this is a pretty nasty hack to ensure the input area viewport doesn't
	// scroll below its content, which can happen when the input viewport
	// height has grown, or previously entered lines have been deleted
	if prevInputLineCnt != m.input.LineCount() {
		// dist is the distance we'd need to scroll up from the current cursor
		// position to position the last input line at the bottom of the
		// viewport. if negative, we're already scrolled above the bottom
		dist := m.input.Line() - (m.input.LineCount() - m.input.Height())
		if dist > 0 {
			for i := 0; i < dist; i++ {
				// move cursor up until content reaches the bottom of the viewport
				m.input.CursorUp()
			}
			m.input, cmd = m.input.Update(nil)
			for i := 0; i < dist; i++ {
				// move cursor back down to its previous position
				m.input.CursorDown()
			}
			m.input, cmd = m.input.Update(nil)
		}
	}

	return m, tea.Batch(cmds...)
}

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
}

type CycleDirection int

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

func cycleMessages(m *models.Message, msgs []models.Message, dir CycleDirection) (*models.Message, error) {
	currentIndex := -1
	for i, reply := range msgs {
		if reply.ID == m.ID {
			currentIndex = i
			break
		}
	}

	if currentIndex < 0 {
		return nil, fmt.Errorf("message not found")
	}

	var next int
	if dir == CyclePrev {
		// Wrap around to the last reply if at the beginning
		next = (currentIndex - 1 + len(msgs)) % len(msgs)
	} else {
		// Wrap around to the first reply if at the end
		next = (currentIndex + 1) % len(msgs)
	}
	return &msgs[next], nil
}

func (m *Model) cycleSelectedRoot(conv *models.Conversation, dir CycleDirection) (*models.Message, error) {
	if len(m.rootMessages) < 2 {
		return nil, nil
	}

	nextRoot, err := cycleMessages(conv.SelectedRoot, m.rootMessages, dir)
	if err != nil {
		return nil, err
	}

	conv.SelectedRoot = nextRoot
	err = m.State.Ctx.Store.UpdateConversation(conv)
	if err != nil {
		return nil, fmt.Errorf("Could not update conversation: %v", err)
	}
	return nextRoot, nil
}

func (m *Model) cycleSelectedReply(message *models.Message, dir CycleDirection) (*models.Message, error) {
	if len(message.Replies) < 2 {
		return nil, nil
	}

	nextReply, err := cycleMessages(message.SelectedReply, message.Replies, dir)
	if err != nil {
		return nil, err
	}

	message.SelectedReply = nextReply
	err = m.State.Ctx.Store.UpdateMessage(message)
	if err != nil {
		return nil, fmt.Errorf("Could not update message: %v", err)
	}
	return nextReply, nil
}

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":
		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
}

func (m *Model) renderMessageHeading(i int, message *models.Message) string {
	icon := ""
	friendly := message.Role.FriendlyRole()
	style := lipgloss.NewStyle().Faint(true).Bold(true)

	switch message.Role {
	case models.MessageRoleSystem:
		icon = "⚙️"
	case models.MessageRoleUser:
		style = userStyle
	case models.MessageRoleAssistant:
		style = assistantStyle
	case models.MessageRoleToolCall:
		style = assistantStyle
		friendly = models.MessageRoleAssistant.FriendlyRole()
	case models.MessageRoleToolResult:
		icon = "🔧"
	}

	user := style.Render(icon + friendly)

	var prefix string
	var suffix string

	faint := lipgloss.NewStyle().Faint(true)

	if i == 0 && len(m.rootMessages) > 0 {
		selectedRootIndex := 0
		for j, reply := range m.rootMessages {
			if reply.ID == *m.conversation.SelectedRootID {
				selectedRootIndex = j
				break
			}
		}
		suffix += faint.Render(fmt.Sprintf(" <%d/%d>", selectedRootIndex+1, len(m.rootMessages)))
	}
	if i > 0 && len(m.messages[i-1].Replies) > 1 {
		// Find the selected reply index
		selectedReplyIndex := 0
		for j, reply := range m.messages[i-1].Replies {
			if reply.ID == *m.messages[i-1].SelectedReplyID {
				selectedReplyIndex = j
				break
			}
		}
		suffix += faint.Render(fmt.Sprintf(" <%d/%d>", selectedReplyIndex+1, len(m.messages[i-1].Replies)))
	}

	if m.focus == focusMessages {
		if i == m.selectedMessage {
			prefix = "> "
		}
	}

	if message.ID == 0 {
		suffix += faint.Render(" (not saved)")
	}

	return messageHeadingStyle.Render(prefix + user + suffix)
}

func (m *Model) renderMessage(i int) string {
	msg := &m.messages[i]

	// Write message contents
	sb := &strings.Builder{}
	sb.Grow(len(msg.Content) * 2)
	if msg.Content != "" {
		err := m.State.Ctx.Chroma.Highlight(sb, msg.Content)
		if err != nil {
			sb.Reset()
			sb.WriteString(msg.Content)
		}
	}

	// Show the assistant's cursor
	if m.waitingForReply && i == len(m.messages)-1 {
		sb.WriteString(m.replyCursor.View())
	}

	// Write tool call info
	var toolString string
	switch msg.Role {
	case models.MessageRoleToolCall:
		bytes, err := yaml.Marshal(msg.ToolCalls)
		if err != nil {
			toolString = "Could not serialize ToolCalls"
		} else {
			toolString = "tool_calls:\n" + string(bytes)
		}
	case models.MessageRoleToolResult:
		if !m.showToolResults {
			break
		}

		type renderedResult struct {
			ToolName string `yaml:"tool"`
			Result   any
		}

		var toolResults []renderedResult
		for _, result := range msg.ToolResults {
			var jsonResult interface{}
			err := json.Unmarshal([]byte(result.Result), &jsonResult)
			if err != nil {
				// If parsing as JSON fails, treat Result as a plain string
				toolResults = append(toolResults, renderedResult{
					ToolName: result.ToolName,
					Result:   result.Result,
				})
			} else {
				// If parsing as JSON succeeds, marshal the parsed JSON into YAML
				toolResults = append(toolResults, renderedResult{
					ToolName: result.ToolName,
					Result:   &jsonResult,
				})
			}
		}

		bytes, err := yaml.Marshal(toolResults)
		if err != nil {
			toolString = "Could not serialize ToolResults"
		} else {
			toolString = "tool_results:\n" + string(bytes)
		}
	}

	if toolString != "" {
		toolString = strings.TrimRight(toolString, "\n")
		if msg.Content != "" {
			sb.WriteString("\n\n")
		}
		_ = m.State.Ctx.Chroma.HighlightLang(sb, toolString, "yaml")
	}

	content := strings.TrimRight(sb.String(), "\n")

	if m.wrap {
		wrapWidth := m.content.Width - messageStyle.GetHorizontalPadding()
		// first we word-wrap text to slightly less than desired width (since
		// wordwrap seems to have an off-by-1 issue), then hard wrap at
		// desired with
		content = wrap.String(wordwrap.String(content, wrapWidth-2), wrapWidth)
	}

	return messageStyle.Width(0).Render(content)
}

// render the conversation into a string
func (m *Model) conversationMessagesView() string {
	sb := strings.Builder{}

	m.messageOffsets = make([]int, len(m.messages))
	lineCnt := 1
	for i, message := range m.messages {
		m.messageOffsets[i] = lineCnt

		switch message.Role {
		case models.MessageRoleToolCall:
			if !m.showToolResults && message.Content == "" {
				continue
			}
		case models.MessageRoleToolResult:
			if !m.showToolResults {
				continue
			}
		}

		heading := m.renderMessageHeading(i, &message)
		sb.WriteString(heading)
		sb.WriteString("\n")
		lineCnt += lipgloss.Height(heading)

		var rendered string
		if m.waitingForReply && i == len(m.messages)-1 {
			// do a direct render of final (assistant) message to handle the
			// assistant cursor blink
			rendered = m.renderMessage(i)
		} else {
			rendered = m.messageCache[i]
		}

		sb.WriteString(rendered)
		sb.WriteString("\n")
		lineCnt += lipgloss.Height(rendered)
	}

	return sb.String()
}

func (m *Model) headerView() string {
	titleStyle := lipgloss.NewStyle().Bold(true)
	var title string
	if m.conversation != nil && m.conversation.Title != "" {
		title = m.conversation.Title
	} else {
		title = "Untitled"
	}
	title = tuiutil.TruncateToCellWidth(title, m.width-styles.Header.GetHorizontalPadding(), "...")
	header := titleStyle.Render(title)
	return styles.Header.Width(m.width).Render(header)
}

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{}

	if m.elapsed > 0 && m.tokenCount > 0 {
		throughput := fmt.Sprintf("%.0f t/sec", float64(m.tokenCount)/m.elapsed.Seconds())
		rightSegments = append(rightSegments, segmentStyle.Render(throughput))
	}

	model := fmt.Sprintf("Model: %s", *m.State.Ctx.Config.Defaults.Model)
	rightSegments = append(rightSegments, segmentStyle.Render(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 = tuiutil.TruncateToCellWidth(footer, m.width, "...")
	}
	return footerStyle.Width(m.width).Render(footer)
}

func (m *Model) setMessage(i int, msg models.Message) {
	if i >= len(m.messages) {
		panic("i out of range")
	}
	m.messages[i] = msg
	m.messageCache[i] = m.renderMessage(i)
}

func (m *Model) addMessage(msg models.Message) {
	m.messages = append(m.messages, msg)
	m.messageCache = append(m.messageCache, m.renderMessage(len(m.messages)-1))
}

func (m *Model) setMessageContents(i int, content string) {
	if i >= len(m.messages) {
		panic("i out of range")
	}
	m.messages[i].Content = content
	m.messageCache[i] = m.renderMessage(i)
}

func (m *Model) rebuildMessageCache() {
	m.messageCache = make([]string, len(m.messages))
	for i := range m.messages {
		m.messageCache[i] = m.renderMessage(i)
	}
}

func (m *Model) updateContent() {
	atBottom := m.content.AtBottom()
	m.content.SetContent(m.conversationMessagesView())
	if atBottom {
		// if we were at bottom before the update, scroll with the output
		m.content.GotoBottom()
	}
}

func (m *Model) loadConversation(shortname string) tea.Cmd {
	return func() tea.Msg {
		if shortname == "" {
			return nil
		}
		c, err := m.State.Ctx.Store.ConversationByShortName(shortname)
		if err != nil {
			return shared.MsgError(fmt.Errorf("Could not lookup conversation: %v", err))
		}
		if c.ID == 0 {
			return shared.MsgError(fmt.Errorf("Conversation not found: %s", shortname))
		}
		return msgConversationLoaded(c)
	}
}

func (m *Model) loadMessages(c *models.Conversation) tea.Cmd {
	return func() tea.Msg {
		messages, err := m.State.Ctx.Store.PathToLeaf(c.SelectedRoot)
		if err != nil {
			return shared.MsgError(fmt.Errorf("Could not load conversation messages: %v\n", err))
		}
		return msgMessagesLoaded(messages)
	}
}

func (m *Model) persistConversation() error {
	if m.conversation.ID == 0 {
		// Start a new conversation with all messages so far
		c, messages, err := m.State.Ctx.Store.StartConversation(m.messages...)
		if err != nil {
			return err
		}
		m.conversation = c
		m.messages = messages

		return nil
	}

	// else, we'll handle updating an existing conversation's messages
	for i := 0; i < len(m.messages); i++ {
		if m.messages[i].ID > 0 {
			// message has an ID, update its contents
			// TODO: check for content/tool equality before updating?
			err := m.State.Ctx.Store.UpdateMessage(&m.messages[i])
			if err != nil {
				return err
			}
		} else if i > 0 {
			// messages is new, so add it as a reply to previous message
			saved, err := m.State.Ctx.Store.Reply(&m.messages[i-1], m.messages[i])
			if err != nil {
				return err
			}
			// add this message as a reply to the previous
			m.messages[i-1].Replies = append(m.messages[i-1].Replies, saved[0])
			m.messages[i] = saved[0]
		} else {
			// message has no id and no previous messages to add it to
			// this shouldn't happen?
			return fmt.Errorf("Error: no messages to reply to")
		}
	}

	return nil
}

func (m *Model) generateConversationTitle() tea.Cmd {
	return func() tea.Msg {
		title, err := cmdutil.GenerateTitle(m.State.Ctx, m.messages)
		if err != nil {
			return shared.MsgError(err)
		}
		return msgConversationTitleChanged(title)
	}
}

func (m *Model) waitForReply() tea.Cmd {
	return func() tea.Msg {
		return msgAssistantReply(<-m.replyChan)
	}
}

func (m *Model) waitForChunk() tea.Cmd {
	return func() tea.Msg {
		return msgResponseChunk(<-m.replyChunkChan)
	}
}

func (m *Model) promptLLM() tea.Cmd {
	m.waitingForReply = true
	m.replyCursor.Blink = false
	m.status = "Press ctrl+c to cancel"

	toPrompt := m.messages

	// Add response placeholder message
	if m.messages[len(m.messages)-1].Role != models.MessageRoleAssistant {
		m.addMessage(models.Message{
			Role:    models.MessageRoleAssistant,
			Content: "",
		})
	}

	m.tokenCount = 0
	m.startTime = time.Now()
	m.elapsed = 0

	return func() tea.Msg {
		completionProvider, err := m.State.Ctx.GetCompletionProvider(*m.State.Ctx.Config.Defaults.Model)
		if err != nil {
			return shared.MsgError(err)
		}

		requestParams := models.RequestParameters{
			Model:       *m.State.Ctx.Config.Defaults.Model,
			MaxTokens:   *m.State.Ctx.Config.Defaults.MaxTokens,
			Temperature: *m.State.Ctx.Config.Defaults.Temperature,
			ToolBag:     m.State.Ctx.EnabledTools,
		}

		replyHandler := func(msg models.Message) {
			m.replyChan <- msg
		}

		ctx, cancel := context.WithCancel(context.Background())

		canceled := false
		go func() {
			select {
			case <-m.stopSignal:
				canceled = true
				cancel()
			}
		}()

		resp, err := completionProvider.CreateChatCompletionStream(
			ctx, requestParams, toPrompt, replyHandler, m.replyChunkChan,
		)

		if err != nil && !canceled {
			return msgResponseError(err)
		}

		return msgResponseEnd(resp)
	}
}