opencode/internal/tui/components/chat/messages.go

package chat

import (
	"fmt"
	"math"
	"time"

	"github.com/charmbracelet/bubbles/key"
	"github.com/charmbracelet/bubbles/spinner"
	"github.com/charmbracelet/bubbles/viewport"
	tea "github.com/charmbracelet/bubbletea"
	"github.com/charmbracelet/lipgloss"
	"github.com/sst/opencode/internal/message"
	"github.com/sst/opencode/internal/session"
	"github.com/sst/opencode/internal/tui/app"
	"github.com/sst/opencode/internal/tui/components/dialog"
	"github.com/sst/opencode/internal/tui/state"
	"github.com/sst/opencode/internal/tui/styles"
	"github.com/sst/opencode/internal/tui/theme"
	"github.com/sst/opencode/pkg/client"
)

type cacheItem struct {
	width   int
	content []uiMessage
}

type messagesCmp struct {
	app              *app.App
	width, height    int
	viewport         viewport.Model
	uiMessages       []uiMessage
	currentMsgID     string
	cachedContent    map[string]cacheItem
	spinner          spinner.Model
	rendering        bool
	attachments      viewport.Model
	showToolMessages bool
}
type renderFinishedMsg struct{}
type ToggleToolMessagesMsg struct{}

type MessageKeys struct {
	PageDown     key.Binding
	PageUp       key.Binding
	HalfPageUp   key.Binding
	HalfPageDown key.Binding
}

var messageKeys = MessageKeys{
	PageDown: key.NewBinding(
		key.WithKeys("pgdown"),
		key.WithHelp("f/pgdn", "page down"),
	),
	PageUp: key.NewBinding(
		key.WithKeys("pgup"),
		key.WithHelp("b/pgup", "page up"),
	),
	HalfPageUp: key.NewBinding(
		key.WithKeys("ctrl+u"),
		key.WithHelp("ctrl+u", "½ page up"),
	),
	HalfPageDown: key.NewBinding(
		key.WithKeys("ctrl+d", "ctrl+d"),
		key.WithHelp("ctrl+d", "½ page down"),
	),
}

func (m *messagesCmp) Init() tea.Cmd {
	return tea.Batch(m.viewport.Init(), m.spinner.Tick)
}

func (m *messagesCmp) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
	m.renderView()

	var cmds []tea.Cmd
	switch msg := msg.(type) {
	case dialog.ThemeChangedMsg:
		m.rerender()
		return m, nil
	case ToggleToolMessagesMsg:
		m.showToolMessages = !m.showToolMessages
		// Clear the cache to force re-rendering of all messages
		m.cachedContent = make(map[string]cacheItem)
		m.renderView()
		return m, nil
	case state.SessionSelectedMsg:
		cmd := m.Reload(msg)
		return m, cmd
	case state.SessionClearedMsg:
		// m.messages = make([]message.Message, 0)
		m.currentMsgID = ""
		m.rendering = false
		return m, nil
	case tea.KeyMsg:
		if key.Matches(msg, messageKeys.PageUp) || key.Matches(msg, messageKeys.PageDown) ||
			key.Matches(msg, messageKeys.HalfPageUp) || key.Matches(msg, messageKeys.HalfPageDown) {
			u, cmd := m.viewport.Update(msg)
			m.viewport = u
			cmds = append(cmds, cmd)
		}
	case renderFinishedMsg:
		m.rendering = false
		m.viewport.GotoBottom()

	case state.StateUpdatedMsg:
		m.renderView()
		m.viewport.GotoBottom()

		// case pubsub.Event[message.Message]:
		// 	needsRerender := false
		// 	if msg.Type == message.EventMessageCreated {
		// 		if msg.Payload.SessionID == m.app.CurrentSessionOLD.ID {
		// 			messageExists := false
		// 			for _, v := range m.messages {
		// 				if v.ID == msg.Payload.ID {
		// 					messageExists = true
		// 					break
		// 				}
		// 			}
		//
		// 			if !messageExists {
		// 				if len(m.messages) > 0 {
		// 					lastMsgID := m.messages[len(m.messages)-1].ID
		// 					delete(m.cachedContent, lastMsgID)
		// 				}
		//
		// 				m.messages = append(m.messages, msg.Payload)
		// 				delete(m.cachedContent, m.currentMsgID)
		// 				m.currentMsgID = msg.Payload.ID
		// 				needsRerender = true
		// 			}
		// 		}
		// 		// There are tool calls from the child task
		// 		for _, v := range m.messages {
		// 			for _, c := range v.ToolCalls() {
		// 				if c.ID == msg.Payload.SessionID {
		// 					delete(m.cachedContent, v.ID)
		// 					needsRerender = true
		// 				}
		// 			}
		// 		}
		// 	} else if msg.Type == message.EventMessageUpdated && msg.Payload.SessionID == m.app.CurrentSessionOLD.ID {
		// 		for i, v := range m.messages {
		// 			if v.ID == msg.Payload.ID {
		// 				m.messages[i] = msg.Payload
		// 				delete(m.cachedContent, msg.Payload.ID)
		// 				needsRerender = true
		// 				break
		// 			}
		// 		}
		// 	}
		// 	if needsRerender {
		// 		m.renderView()
		// 		if len(m.messages) > 0 {
		// 			if (msg.Type == message.EventMessageCreated) ||
		// 				(msg.Type == message.EventMessageUpdated && msg.Payload.ID == m.messages[len(m.messages)-1].ID) {
		// 				m.viewport.GotoBottom()
		// 			}
		// 		}
		// 	}
	}

	spinner, cmd := m.spinner.Update(msg)
	m.spinner = spinner
	cmds = append(cmds, cmd)
	return m, tea.Batch(cmds...)
}

func (m *messagesCmp) IsAgentWorking() bool {
	return m.app.PrimaryAgentOLD.IsSessionBusy(m.app.CurrentSessionOLD.ID)
}

func formatTimeDifference(unixTime1, unixTime2 int64) string {
	diffSeconds := float64(math.Abs(float64(unixTime2 - unixTime1)))

	if diffSeconds < 60 {
		return fmt.Sprintf("%.1fs", diffSeconds)
	}

	minutes := int(diffSeconds / 60)
	seconds := int(diffSeconds) % 60
	return fmt.Sprintf("%dm%ds", minutes, seconds)
}

func (m *messagesCmp) renderView() {
	m.uiMessages = make([]uiMessage, 0)
	baseStyle := styles.BaseStyle()

	if m.width == 0 {
		return
	}

	for _, msg := range m.app.Messages {
		switch msg.Role {
		case client.User:
			if cache, ok := m.cachedContent[msg.Id]; ok && cache.width == m.width {
				m.uiMessages = append(m.uiMessages, cache.content...)
				continue
			}
			userMsg := renderUserMessage(
				msg,
				msg.Id == m.currentMsgID,
				m.width,
			)
			m.uiMessages = append(m.uiMessages, userMsg)
			m.cachedContent[msg.Id] = cacheItem{
				width:   m.width,
				content: []uiMessage{userMsg},
			}
			// pos += userMsg.height + 1 // + 1 for spacing
		case client.Assistant:
			if cache, ok := m.cachedContent[msg.Id]; ok && cache.width == m.width {
				m.uiMessages = append(m.uiMessages, cache.content...)
				continue
			}
			// assistantMessages := renderAssistantMessage(
			// 	msg,
			// 	inx,
			// 	m.app.Messages,
			// 	m.app.MessagesOLD,
			// 	m.currentMsgID,
			// 	m.width,
			// 	pos,
			// 	m.showToolMessages,
			// )
			// for _, msg := range assistantMessages {
			// 	m.uiMessages = append(m.uiMessages, msg)
			// 	// pos += msg.height + 1 // + 1 for spacing
			// }
			// m.cachedContent[msg.Id] = cacheItem{
			// 	width:   m.width,
			// 	content: assistantMessages,
			// }
		}
	}

	messages := make([]string, 0)
	for _, v := range m.uiMessages {
		messages = append(messages, lipgloss.JoinVertical(lipgloss.Left, v.content),
			baseStyle.
				Width(m.width).
				Render(
					"",
				),
		)
	}

	// temp, _ := json.MarshalIndent(m.app.State, "", "    ")

	m.viewport.SetContent(
		baseStyle.
			Width(m.width).
			Render(
				// string(temp),
				lipgloss.JoinVertical(
					lipgloss.Top,
					messages...,
				),
			),
	)
}

func (m *messagesCmp) View() string {
	baseStyle := styles.BaseStyle()

	if m.rendering {
		return baseStyle.
			Width(m.width).
			Render(
				lipgloss.JoinVertical(
					lipgloss.Top,
					"Loading...",
					// m.working(),
					m.help(),
				),
			)
	}
	if len(m.app.Messages) == 0 {
		content := baseStyle.
			Width(m.width).
			Height(m.height - 1).
			Render(
				m.initialScreen(),
			)

		return baseStyle.
			Width(m.width).
			Render(
				lipgloss.JoinVertical(
					lipgloss.Top,
					content,
					"",
					m.help(),
				),
			)
	}

	return baseStyle.
		Width(m.width).
		Render(
			lipgloss.JoinVertical(
				lipgloss.Top,
				m.viewport.View(),
				// m.working(),
				m.help(),
			),
		)
}

func hasToolsWithoutResponse(messages []message.Message) bool {
	toolCalls := make([]message.ToolCall, 0)
	toolResults := make([]message.ToolResult, 0)
	for _, m := range messages {
		toolCalls = append(toolCalls, m.ToolCalls()...)
		toolResults = append(toolResults, m.ToolResults()...)
	}

	for _, v := range toolCalls {
		found := false
		for _, r := range toolResults {
			if v.ID == r.ToolCallID {
				found = true
				break
			}
		}
		if !found && v.Finished {
			return true
		}
	}
	return false
}

func hasUnfinishedToolCalls(messages []message.Message) bool {
	toolCalls := make([]message.ToolCall, 0)
	for _, m := range messages {
		toolCalls = append(toolCalls, m.ToolCalls()...)
	}
	for _, v := range toolCalls {
		if !v.Finished {
			return true
		}
	}
	return false
}

// func (m *messagesCmp) working() string {
// 	text := ""
// 	if m.IsAgentWorking() && len(m.app.Messages) > 0 {
// 		t := theme.CurrentTheme()
// 		baseStyle := styles.BaseStyle()
//
// 		task := "Thinking..."
// 		lastMessage := m.app.Messages[len(m.app.Messages)-1]
// 		if hasToolsWithoutResponse(m.app.Messages) {
// 			task = "Waiting for tool response..."
// 		} else if hasUnfinishedToolCalls(m.app.Messages) {
// 			task = "Building tool call..."
// 		} else if !lastMessage.IsFinished() {
// 			task = "Generating..."
// 		}
// 		if task != "" {
// 			text += baseStyle.
// 				Width(m.width).
// 				Foreground(t.Primary()).
// 				Bold(true).
// 				Render(fmt.Sprintf("%s %s ", m.spinner.View(), task))
// 		}
// 	}
// 	return text
// }

func (m *messagesCmp) help() string {
	t := theme.CurrentTheme()
	baseStyle := styles.BaseStyle()

	text := ""

	if m.app.PrimaryAgentOLD.IsBusy() {
		text += lipgloss.JoinHorizontal(
			lipgloss.Left,
			baseStyle.Foreground(t.TextMuted()).Bold(true).Render("press "),
			baseStyle.Foreground(t.Text()).Bold(true).Render("esc"),
			baseStyle.Foreground(t.TextMuted()).Bold(true).Render(" to interrupt"),
		)
	} else {
		text += lipgloss.JoinHorizontal(
			lipgloss.Left,
			baseStyle.Foreground(t.Text()).Bold(true).Render("enter"),
			baseStyle.Foreground(t.TextMuted()).Bold(true).Render(" to send,"),
			baseStyle.Foreground(t.Text()).Bold(true).Render(" \\"),
			baseStyle.Foreground(t.TextMuted()).Bold(true).Render("+"),
			baseStyle.Foreground(t.Text()).Bold(true).Render("enter"),
			baseStyle.Foreground(t.TextMuted()).Bold(true).Render(" for newline,"),
			baseStyle.Foreground(t.Text()).Bold(true).Render(" ↑↓"),
			baseStyle.Foreground(t.TextMuted()).Bold(true).Render(" for history,"),
			baseStyle.Foreground(t.Text()).Bold(true).Render(" ctrl+h"),
			baseStyle.Foreground(t.TextMuted()).Bold(true).Render(" to toggle tool messages"),
		)
	}
	return baseStyle.
		Width(m.width).
		Render(text)
}

func (m *messagesCmp) initialScreen() string {
	baseStyle := styles.BaseStyle()

	return baseStyle.Width(m.width).Render(
		lipgloss.JoinVertical(
			lipgloss.Top,
			header(m.width),
			"",
			lspsConfigured(m.width),
		),
	)
}

func (m *messagesCmp) rerender() {
	for _, msg := range m.app.Messages {
		delete(m.cachedContent, msg.Id)
	}
	m.renderView()
}

func (m *messagesCmp) SetSize(width, height int) tea.Cmd {
	if m.width == width && m.height == height {
		return nil
	}
	m.width = width
	m.height = height
	m.viewport.Width = width
	m.viewport.Height = height - 2
	m.attachments.Width = width + 40
	m.attachments.Height = 3
	m.rerender()
	return nil
}

func (m *messagesCmp) GetSize() (int, int) {
	return m.width, m.height
}

func (m *messagesCmp) Reload(session *session.Session) tea.Cmd {
	// messages := m.app.Messages
	// messages, err := m.app.MessagesOLD.List(context.Background(), session.ID)
	// if err != nil {
	// 	status.Error(err.Error())
	// 	return nil
	// }
	// m.messages = messages

	if len(m.app.Messages) > 0 {
		m.currentMsgID = m.app.Messages[len(m.app.Messages)-1].Id
	}
	delete(m.cachedContent, m.currentMsgID)
	m.rendering = true
	return func() tea.Msg {
		m.renderView()
		return renderFinishedMsg{}
	}
}

func (m *messagesCmp) BindingKeys() []key.Binding {
	return []key.Binding{
		m.viewport.KeyMap.PageDown,
		m.viewport.KeyMap.PageUp,
		m.viewport.KeyMap.HalfPageUp,
		m.viewport.KeyMap.HalfPageDown,
	}
}

func NewMessagesCmp(app *app.App) tea.Model {
	customSpinner := spinner.Spinner{
		Frames: []string{" ", "┃", "┃"},
		FPS:    time.Second / 3,
	}
	s := spinner.New(spinner.WithSpinner(customSpinner))
	vp := viewport.New(0, 0)
	attachmets := viewport.New(0, 0)
	vp.KeyMap.PageUp = messageKeys.PageUp
	vp.KeyMap.PageDown = messageKeys.PageDown
	vp.KeyMap.HalfPageUp = messageKeys.HalfPageUp
	vp.KeyMap.HalfPageDown = messageKeys.HalfPageDown
	return &messagesCmp{
		app:              app,
		cachedContent:    make(map[string]cacheItem),
		viewport:         vp,
		spinner:          s,
		attachments:      attachmets,
		showToolMessages: true,
	}
}