agentu/internal/tui/app.go

package tui

import (
	"context"
	"errors"
	"fmt"
	"io"
	"math"
	"strings"

	"agentu/internal/agent"

	"github.com/charmbracelet/bubbles/textarea"
	"github.com/charmbracelet/bubbles/viewport"
	tea "github.com/charmbracelet/bubbletea"
	"github.com/charmbracelet/lipgloss"
)

const (
	minInputLines = 1
	maxInputLines = 6
)

type Options struct {
	ModelName    string
	Yolo         bool
	ThemeMode    ThemeMode
	ModelManager ModelManager
}

type ModelManager interface {
	CurrentProvider() string
	CurrentModel() string
	CurrentThinking() string
	Info() string
	SetProvider(name string) (string, error)
	SetModel(model string) (string, error)
	SetThinking(level string) (string, error)
}

func Run(ctx context.Context, assistant *agent.Agent, opts Options) error {
	model := newModel(ctx, assistant, opts)
	_, err := tea.NewProgram(model, tea.WithAltScreen()).Run()
	return err
}

type role string

const (
	roleUser      role = "user"
	roleAssistant role = "assistant"
	roleTool      role = "tool"
	roleError     role = "error"
	roleSystem    role = "system"
)

type message struct {
	role    role
	content string
}

type model struct {
	ctx       context.Context
	agent     *agent.Agent
	modelName string
	yolo      bool
	models    ModelManager
	theme     theme
	styles    styles

	viewport viewport.Model
	input    textarea.Model

	width  int
	height int

	messages []message
	status   string
	running  bool
	cancel   context.CancelFunc
	events   <-chan tea.Msg
}

type assistantChunkMsg string
type toolLogMsg string
type turnStartedMsg struct{}
type turnDoneMsg struct {
	err error
}

type slashCommand struct {
	Name        string
	Description string
}

var slashCommands = []slashCommand{
	{Name: "/clear", Description: "reset context"},
	{Name: "/exit", Description: "quit"},
	{Name: "/quit", Description: "quit"},
	{Name: "/model", Description: "model/provider/thinking"},
}

func newModel(ctx context.Context, assistant *agent.Agent, opts Options) model {
	mode := opts.ThemeMode
	if mode == "" {
		mode = ThemeLight
	}
	th := themeForMode(mode)
	st := th.styles()

	input := textarea.New()
	input.Placeholder = "Message agentu..."
	input.Prompt = ""
	input.ShowLineNumbers = false
	input.MaxHeight = maxInputLines
	input.SetHeight(minInputLines)
	input.SetWidth(80)
	applyTextareaTheme(&input, th)
	input.Focus()

	vp := viewport.New(80, 20)
	vp.Style = st.Viewport
	vp.SetContent("")

	return model{
		ctx:       ctx,
		agent:     assistant,
		modelName: opts.ModelName,
		yolo:      opts.Yolo,
		models:    opts.ModelManager,
		theme:     th,
		styles:    st,
		viewport:  vp,
		input:     input,
		status:    "ready",
		messages: []message{
			{role: roleSystem, content: "Ask anything. Enter sends, ctrl+j or alt+enter inserts a newline, esc cancels a running turn."},
		},
	}
}

func (m model) Init() tea.Cmd {
	return textarea.Blink
}

func (m model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
	var cmds []tea.Cmd

	switch msg := msg.(type) {
	case tea.WindowSizeMsg:
		m.width = msg.Width
		m.height = msg.Height
		m.resize()
		m.refreshViewport(true)
		return m, nil

	case tea.KeyMsg:
		switch msg.String() {
		case "ctrl+c":
			if m.cancel != nil {
				m.cancel()
			}
			return m, tea.Quit
		case "esc":
			if m.running && m.cancel != nil {
				m.cancel()
				m.status = "canceling..."
				return m, nil
			}
		case "enter":
			if m.running {
				return m, nil
			}
			return m.submit()
		case "alt+enter", "ctrl+j":
			return m.insertInputNewline()
		}

	case assistantChunkMsg:
		m.appendAssistantChunk(string(msg))
		m.status = "streaming..."
		m.refreshViewport(true)
		return m, m.waitForEvent()

	case toolLogMsg:
		text := cleanToolLog(string(msg))
		if text != "" {
			m.messages = append(m.messages, message{role: roleTool, content: text})
			m.status = text
			m.refreshViewport(true)
		}
		return m, m.waitForEvent()

	case turnDoneMsg:
		m.running = false
		m.events = nil
		m.cancel = nil
		if msg.err != nil {
			m.messages = append(m.messages, message{role: roleError, content: msg.err.Error()})
			m.status = "error"
		} else {
			m.status = "ready"
		}
		m.refreshViewport(true)
		cmds = append(cmds, textarea.Blink)
	}

	if keyMsg, ok := msg.(tea.KeyMsg); ok {
		if isViewportKey(keyMsg.String()) {
			nextViewport, cmd := m.viewport.Update(msg)
			m.viewport = nextViewport
			return m, cmd
		}
		if !m.running {
			nextInput, cmd := m.input.Update(msg)
			m.input = nextInput
			m.afterInputChanged()
			return m, cmd
		}
		return m, nil
	}

	if !m.running {
		nextInput, cmd := m.input.Update(msg)
		m.input = nextInput
		m.afterInputChanged()
		cmds = append(cmds, cmd)
	}

	nextViewport, cmd := m.viewport.Update(msg)
	m.viewport = nextViewport
	cmds = append(cmds, cmd)

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

func (m model) View() string {
	if m.width <= 0 || m.height <= 0 {
		return "agentu"
	}

	header := m.headerView()
	status := m.statusView()
	input := m.inputView()
	body := lipgloss.JoinVertical(lipgloss.Left, header, m.viewport.View(), input, status)
	return m.styles.App.Width(m.width).Height(m.height).Render(body)
}

func (m *model) submit() (tea.Model, tea.Cmd) {
	input := strings.TrimSpace(m.input.Value())
	if input == "" {
		return *m, nil
	}

	switch input {
	case "/exit", "/quit":
		return *m, tea.Quit
	case "/clear":
		m.agent.Clear()
		m.messages = []message{{role: roleSystem, content: "Context cleared."}}
		m.input.SetValue("")
		m.afterInputChanged()
		m.status = "ready"
		m.refreshViewport(true)
		return *m, nil
	case "/model":
		return m.runLocalCommand(input)
	}

	if strings.HasPrefix(input, "/") {
		return m.runLocalCommand(input)
	}

	m.messages = append(m.messages, message{role: roleUser, content: input})
	m.input.SetValue("")
	m.afterInputChanged()
	m.running = true
	m.status = "thinking..."
	m.refreshViewport(true)

	turnCtx, cancel := context.WithCancel(m.ctx)
	m.cancel = cancel

	events := make(chan tea.Msg, 256)
	m.events = events

	start := func() tea.Msg {
		go func() {
			defer close(events)
			out := eventWriter{ctx: turnCtx, events: events, assistant: true}
			logs := eventWriter{ctx: turnCtx, events: events}
			err := m.agent.RunTurn(turnCtx, input, out, logs)
			if errors.Is(err, context.Canceled) {
				err = nil
			}
			select {
			case events <- turnDoneMsg{err: err}:
			case <-turnCtx.Done():
			}
		}()
		return turnStartedMsg{}
	}

	return *m, tea.Batch(start, m.waitForEvent())
}

func (m *model) insertInputNewline() (tea.Model, tea.Cmd) {
	enter := tea.KeyMsg{Type: tea.KeyEnter}
	nextInput, cmd := m.input.Update(enter)
	m.input = nextInput
	m.afterInputChanged()
	return *m, cmd
}

func (m model) waitForEvent() tea.Cmd {
	events := m.events
	if events == nil {
		return nil
	}
	return func() tea.Msg {
		msg, ok := <-events
		if !ok {
			return turnDoneMsg{}
		}
		return msg
	}
}

func (m *model) appendAssistantChunk(chunk string) {
	if chunk == "" {
		return
	}
	if len(m.messages) == 0 || m.messages[len(m.messages)-1].role != roleAssistant {
		m.messages = append(m.messages, message{role: roleAssistant})
	}
	m.messages[len(m.messages)-1].content += chunk
}

func (m *model) resize() {
	width := max(40, m.width)
	inputWidth := max(20, width-4)
	m.input.SetWidth(inputWidth)
	m.syncInputHeight()

	headerHeight := 1
	statusHeight := 1
	inputHeight := m.inputBlockHeight()
	m.viewport.Width = width
	m.viewport.Height = max(4, m.height-headerHeight-statusHeight-inputHeight)
	m.viewport.Style = m.styles.Viewport
}

func (m *model) refreshViewport(bottom bool) {
	m.viewport.SetContent(m.renderMessages())
	if bottom {
		m.viewport.GotoBottom()
	}
}

func (m model) renderMessages() string {
	width := max(40, m.width)
	contentWidth := max(20, width-6)
	parts := make([]string, 0, len(m.messages))
	for _, msg := range m.messages {
		content := strings.TrimRight(msg.content, "\n")
		if content == "" {
			continue
		}
		parts = append(parts, m.messageView(msg.role, content, contentWidth))
	}
	if len(parts) == 0 {
		return m.styles.Muted.Width(contentWidth).Render("Start a conversation.")
	}
	return strings.Join(parts, "\n\n")
}

func (m model) messageView(r role, content string, width int) string {
	label := string(r)
	style := m.styles.AssistantMsg
	labelStyle := m.styles.AssistantLabel

	switch r {
	case roleUser:
		label = "you"
		style = m.styles.UserMessage
		labelStyle = m.styles.UserLabel
	case roleTool:
		label = "tool"
		style = m.styles.ToolMessage
		labelStyle = m.styles.ToolLabel
	case roleError:
		label = "error"
		style = m.styles.ErrorMessage
		labelStyle = m.styles.ErrorLabel
	case roleSystem:
		label = "note"
		style = m.styles.SystemMessage
		labelStyle = m.styles.SystemLabel
	default:
		label = "agentu"
	}

	labelLine := labelStyle.Render(label)
	body := style.Width(width).Render(content)
	return lipgloss.JoinVertical(lipgloss.Left, labelLine, body)
}

func (m model) headerView() string {
	mode := "read-only tools"
	if m.yolo {
		mode = "yolo tools"
	}
	modelName := m.modelName
	thinking := ""
	if m.models != nil {
		modelName = m.models.CurrentModel()
		if currentThinking := m.models.CurrentThinking(); currentThinking != "" && currentThinking != "unset" {
			thinking = " · thinking " + currentThinking
		}
	}
	title := m.styles.Title.Render("agentu")
	meta := m.styles.Muted.Render(fmt.Sprintf("%s · %s%s · %s", modelName, mode, thinking, m.theme.Mode))
	line := lipgloss.JoinHorizontal(lipgloss.Center, title, "  ", meta)
	return m.styles.Header.Width(max(1, m.width)).Render(line)
}

func (m model) inputView() string {
	title := m.styles.Muted.Render("Enter sends · ctrl+j newline · alt+enter newline · /clear · /exit")
	if m.running {
		title = m.styles.Muted.Render("Working · esc cancel · ctrl+c quit")
	}
	suggestions := m.slashSuggestionsView()
	box := m.styles.InputBox.Width(max(20, m.width-2)).Render(m.input.View())
	if suggestions != "" {
		return lipgloss.JoinVertical(lipgloss.Left, suggestions, title, box)
	}
	return lipgloss.JoinVertical(lipgloss.Left, title, box)
}

func (m model) statusView() string {
	left := m.status
	if left == "" {
		left = "ready"
	}
	return m.styles.Status.Width(max(1, m.width)).Render(left)
}

type eventWriter struct {
	ctx       context.Context
	events    chan<- tea.Msg
	assistant bool
}

func (w eventWriter) Write(p []byte) (int, error) {
	if len(p) == 0 {
		return 0, nil
	}
	text := string(append([]byte(nil), p...))
	var msg tea.Msg
	if w.assistant {
		msg = assistantChunkMsg(text)
	} else {
		msg = toolLogMsg(text)
	}
	select {
	case w.events <- msg:
		return len(p), nil
	case <-w.ctx.Done():
		return 0, w.ctx.Err()
	}
}

func cleanToolLog(text string) string {
	text = strings.TrimSpace(text)
	text = strings.TrimPrefix(text, "[tool]")
	return strings.TrimSpace(text)
}

func isViewportKey(key string) bool {
	switch key {
	case "pgup", "pgdown", "ctrl+up", "ctrl+down":
		return true
	default:
		return false
	}
}

func (m *model) afterInputChanged() {
	m.syncInputHeight()
	if m.width > 0 && m.height > 0 {
		m.resize()
		m.refreshViewport(true)
	}
}

func (m *model) syncInputHeight() {
	m.input.SetHeight(m.desiredInputHeight())
}

func (m model) desiredInputHeight() int {
	value := m.input.Value()
	if value == "" {
		return minInputLines
	}

	textWidth := max(1, m.input.Width())
	lines := 0
	for _, line := range strings.Split(value, "\n") {
		displayWidth := lipgloss.Width(line)
		wrapped := int(math.Ceil(float64(max(1, displayWidth)) / float64(textWidth)))
		lines += max(1, wrapped)
	}
	return min(max(lines, minInputLines), maxInputLines)
}

func (m model) inputBlockHeight() int {
	height := 1 + m.input.Height() + m.styles.InputBox.GetVerticalFrameSize()
	if m.showSlashSuggestions() {
		height++
	}
	return height
}

func (m model) showSlashSuggestions() bool {
	value := strings.TrimSpace(m.input.Value())
	return strings.HasPrefix(value, "/") && !strings.Contains(value, " ")
}

func (m model) slashSuggestionsView() string {
	if !m.showSlashSuggestions() {
		return ""
	}
	prefix := strings.TrimSpace(m.input.Value())
	var parts []string
	for _, command := range slashCommands {
		if !strings.HasPrefix(command.Name, prefix) {
			continue
		}
		parts = append(parts, fmt.Sprintf("%s %s", command.Name, command.Description))
	}
	if len(parts) == 0 {
		parts = append(parts, "No matching commands")
	}
	return m.styles.CommandHint.Width(max(20, m.width-2)).Render(strings.Join(parts, "  "))
}

func (m model) modelInfo() string {
	if m.models != nil {
		return m.models.Info()
	}
	mode := "read-only tools"
	if m.yolo {
		mode = "yolo tools"
	}
	return fmt.Sprintf("model: %s\nmode: %s\ntheme: %s", m.modelName, mode, m.theme.Mode)
}

func (m *model) runLocalCommand(input string) (tea.Model, tea.Cmd) {
	output, err := m.handleLocalCommand(input)
	role := roleSystem
	status := "command"
	if err != nil {
		output = err.Error()
		role = roleError
		status = "command error"
	}
	m.messages = append(m.messages, message{role: role, content: output})
	m.input.SetValue("")
	m.afterInputChanged()
	m.status = status
	m.refreshViewport(true)
	return *m, nil
}

func (m model) handleLocalCommand(input string) (string, error) {
	fields := strings.Fields(input)
	if len(fields) == 0 {
		return "", fmt.Errorf("empty command")
	}
	switch fields[0] {
	case "/model":
		return m.handleModelCommand(fields[1:])
	default:
		return "", fmt.Errorf("unknown command: %s", fields[0])
	}
}

func (m model) handleModelCommand(args []string) (string, error) {
	if m.models == nil {
		if len(args) > 0 {
			return "", fmt.Errorf("model switching is not configured")
		}
		return m.modelInfo(), nil
	}
	if len(args) == 0 || args[0] == "help" {
		return m.models.Info(), nil
	}
	if len(args) == 1 {
		return m.models.SetModel(args[0])
	}
	switch args[0] {
	case "provider":
		if len(args) != 2 {
			return "", fmt.Errorf("usage: /model provider <name>")
		}
		return m.models.SetProvider(args[1])
	case "model":
		if len(args) != 2 {
			return "", fmt.Errorf("usage: /model model <name>")
		}
		return m.models.SetModel(args[1])
	case "thinking":
		if len(args) != 2 {
			return "", fmt.Errorf("usage: /model thinking <none|middle|high|xhigh|max>")
		}
		return m.models.SetThinking(args[1])
	default:
		return "", fmt.Errorf("unknown /model command %q", args[0])
	}
}

var _ io.Writer = eventWriter{}