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

package chat

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

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

// 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) View() string {
	if m.Width == 0 {
		return ""
	}
	sections := make([]string, 0, 6)

	if m.Header != "" {
		sections = append(sections, m.Header)
	}

	sections = append(sections, m.Content)
	if m.Error != "" {
		sections = append(sections, m.Error)
	}
	sections = append(sections, m.Input)

	if m.Footer != "" {
		sections = append(sections, m.Footer)
	}

	return lipgloss.JoinVertical(lipgloss.Left, sections...)
}

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) > 1 && m.conversation.SelectedRootID != nil {
		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.Shared.Ctx.Chroma.Highlight(sb, msg.Content)
		if err != nil {
			sb.Reset()
			sb.WriteString(msg.Content)
		}
	}

	// Show the assistant's cursor
	if m.state == pendingResponse && i == len(m.messages)-1 && msg.Role == models.MessageRoleAssistant {
		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.Shared.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.state == pendingResponse && 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)
	}

	// Render a placeholder for the incoming assistant reply
	if m.state == pendingResponse && (len(m.messages) == 0 || m.messages[len(m.messages)-1].Role != models.MessageRoleAssistant) {
		heading := m.renderMessageHeading(-1, &models.Message{
			Role: models.MessageRoleAssistant,
		})
		sb.WriteString(heading)
		sb.WriteString("\n")
		sb.WriteString(messageStyle.Width(0).Render(m.replyCursor.View()))
		sb.WriteString("\n")
	}

	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("❌💾")
	}

	var status string
	switch m.state {
	case pendingResponse:
		status = "Press ctrl+c to cancel" + m.spinner.View()
	default:
		status = "Press ctrl+s to send"
	}

	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.Shared.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)
}