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

package chat

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

	"git.mlow.ca/mlow/lmcli/pkg/api"
	"git.mlow.ca/mlow/lmcli/pkg/conversation"
	"git.mlow.ca/mlow/lmcli/pkg/tui/styles"
	tuiutil "git.mlow.ca/mlow/lmcli/pkg/tui/util"
	"github.com/charmbracelet/lipgloss"
	"github.com/muesli/reflow/wordwrap"
	"github.com/muesli/reflow/wrap"
	"gopkg.in/yaml.v3"
)

// styles
var (
	boldStyle      = lipgloss.NewStyle().Bold(true)
	faintStyle     = lipgloss.NewStyle().Faint(true)
	boldFaintStyle = lipgloss.NewStyle().Faint(true).Bold(true)

	messageHeadingStyle = lipgloss.NewStyle().
				MarginTop(1).
				MarginBottom(1)

	userStyle = boldFaintStyle.Foreground(lipgloss.Color("10"))

	assistantStyle = boldFaintStyle.Foreground(lipgloss.Color("12"))

	systemStyle = boldStyle.Foreground(lipgloss.Color("8"))

	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().Padding(0, 1)
)

func (m *Model) renderMessageHeading(i int, message *conversation.Message) string {
	friendly := message.Role.FriendlyRole()
	style := systemStyle

	switch message.Role {
	case api.MessageRoleUser:
		style = userStyle
	case api.MessageRoleAssistant:
		style = assistantStyle
	case api.MessageRoleToolCall:
		style = assistantStyle
		friendly = api.MessageRoleAssistant.FriendlyRole()
	case api.MessageRoleSystem:
	case api.MessageRoleToolResult:
	}

	user := style.Render(friendly)

	var prefix, suffix string

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

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

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

	heading := prefix + user + suffix

	if message.Metadata.GenerationModel != nil && m.showDetails {
		heading += faintStyle.Render(
			fmt.Sprintf(" | %s", *message.Metadata.GenerationModel),
		)
	}

	return messageHeadingStyle.Render(heading)
}

// renderMessages renders the message at the given index as it should be shown
// *at this moment* - we render differently depending on the current application
// state (window size, etc, etc).
func (m *Model) renderMessage(i int) string {
	msg := &m.App.Messages[i]

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

	isLast := i == len(m.App.Messages)-1
	isAssistant := msg.Role == api.MessageRoleAssistant

	if m.state == pendingResponse && isLast && isAssistant {
		// Show the assistant's cursor
		sb.WriteString(m.replyCursor.View())
	}

	// Write tool call info
	var toolString string
	switch msg.Role {
	case api.MessageRoleToolCall:
		bytes, err := yaml.Marshal(msg.ToolCalls)
		if err != nil {
			toolString = "Could not serialize ToolCalls"
		} else {
			toolString = "tool_calls:\n" + string(bytes)
		}
	case api.MessageRoleToolResult:
		type renderedResult struct {
			ToolName string `yaml:"tool"`
			Result   any    `yaml:"result,omitempty"`
		}

		var toolResults []renderedResult
		for _, result := range msg.ToolResults {
			if m.showDetails {
				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,
					})
				}
			} else {
				// Only show the tool name when results are hidden
				toolResults = append(toolResults, renderedResult{
					ToolName: result.ToolName,
					Result:   "(hidden, press ctrl+t to view)",
				})
			}
		}

		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.App.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 {
	m.messageOffsets = make([]int, len(m.App.Messages))
	lineCnt := 1

	sb := strings.Builder{}
	for i, message := range m.App.Messages {
		m.messageOffsets[i] = lineCnt

		if !m.showDetails && message.Role.IsSystem() {
			continue
		}

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

		rendered := m.messageCache[i]
		sb.WriteString(rendered)
		sb.WriteString("\n")
		lineCnt += lipgloss.Height(rendered)
	}

	// Render a placeholder for the incoming assistant reply
	if m.state == pendingResponse && m.App.Messages[len(m.App.Messages)-1].Role != api.MessageRoleAssistant {
		heading := m.renderMessageHeading(-1, &conversation.Message{
			Role: api.MessageRoleAssistant,
			Metadata: conversation.MessageMeta{
				GenerationModel: &m.App.Model,
			},
		})
		sb.WriteString(heading)
		sb.WriteString("\n")
		sb.WriteString(messageStyle.Width(0).Render(m.replyCursor.View()))
		sb.WriteString("\n")
	}
	return sb.String()
}

func (m *Model) Content(width, height int) string {
	// calculate clamped input height to accomodate input text
	// minimum 4 lines, maximum half of content area
	inputHeight := max(4, min(height/2, m.input.LineCount()))
	m.input.SetHeight(inputHeight)
	input := m.input.View()

	// remaining height towards content
	m.content.Width, m.content.Height = width, height-tuiutil.Height(input)
	content := m.content.View()
	return lipgloss.JoinVertical(lipgloss.Left, content, input)
}

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

func (m *Model) Footer(width int) string {
	segmentStyle := lipgloss.NewStyle().Faint(true)
	segmentSeparator := segmentStyle.Render(" | ")

	// Left segments
	leftSegments := make([]string, 0, 4)

	if m.state == pendingResponse {
		leftSegments = append(leftSegments, segmentStyle.Render(m.spinner.View()))
	} else {
		leftSegments = append(leftSegments, segmentStyle.Render("∙∙∙"))
	}

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

	// var status string
	// switch m.state {
	// case pendingResponse:
	// 	status = "Press ctrl+c to cancel"
	// default:
	// 	status = "Press ctrl+s to send"
	// }
	// leftSegments = append(leftSegments, segmentStyle.Render(status))

	// Right segments
	rightSegments := make([]string, 0, 8)

	if m.App.Agent != nil {
		rightSegments = append(rightSegments, segmentStyle.Render(m.App.Agent.Name))
	}

	model := segmentStyle.Render(m.App.ActiveModel(lipgloss.NewStyle()))
	rightSegments = append(rightSegments, model)

	savingStyle := segmentStyle.Bold(true)
	saving := ""
	if m.persistence {
		saving = savingStyle.Foreground(lipgloss.Color("2")).Render("💾✅")
	} else {
		saving = savingStyle.Foreground(lipgloss.Color("1")).Render("💾❌")
	}
	rightSegments = append(rightSegments, saving)

	return m.layoutFooter(width, leftSegments, rightSegments, segmentSeparator)
}

func (m *Model) layoutFooter(
	width int,
	leftSegments []string,
	rightSegments []string,
	segmentSeparator string,
) string {
	left := strings.Join(leftSegments, segmentSeparator)
	right := strings.Join(rightSegments, segmentSeparator)

	leftWidth := tuiutil.Width(left)
	rightWidth := tuiutil.Width(right)
	sepWidth := tuiutil.Width(segmentSeparator)
	frameWidth := footerStyle.GetHorizontalFrameSize()

	availableWidth := width - frameWidth - leftWidth - rightWidth

	if availableWidth >= sepWidth {
		// Everything fits
		padding := strings.Repeat(" ", availableWidth)
		return footerStyle.Render(left + padding + right)
	}

	// Inserted between left and right segments when they're being truncated
	div := "..."

	totalAvailableWidth := width - frameWidth
	availableTruncWidth := totalAvailableWidth - len(div)

	minVisibleLength := 3
	if availableTruncWidth < 2*minVisibleLength {
		minVisibleLength = availableTruncWidth / 2
	}

	leftProportion := float64(leftWidth) / float64(leftWidth+rightWidth)

	newLeftWidth := int(max(float64(minVisibleLength), leftProportion*float64(availableTruncWidth)))
	newRightWidth := totalAvailableWidth - newLeftWidth

	truncatedLeft := faintStyle.Render(tuiutil.TruncateRightToCellWidth(left, newLeftWidth, ""))
	truncatedRight := faintStyle.Render(tuiutil.TruncateLeftToCellWidth(right, newRightWidth, "..."))

	return footerStyle.Width(width).Render(truncatedLeft + truncatedRight)
}