orca.ai/internal/tui/model.go

package tui

import (
	"fmt"
	"strings"
	"time"

	"github.com/charmbracelet/bubbles/textarea"
	"github.com/charmbracelet/bubbles/viewport"
	tea "github.com/charmbracelet/bubbletea"
	"github.com/charmbracelet/lipgloss"
	"github.com/orca/orca/pkg/actor"
	"github.com/orca/orca/pkg/kernel"
)

type ChatMessage struct {
	Role    string
	Content string
	Agent   string
}

type Model struct {
	width     int
	height    int
	kernel    *kernel.Kernel
	viewport  viewport.Model
	textarea  textarea.Model
	messages  []ChatMessage
	events    chan Event
	writer    *EventWriter
	ready     bool
	loading   bool
	agentRuns map[string]bool
}

func (m Model) TestEvents() chan Event {
	return m.events
}

func (m Model) MessageCount() int {
	return len(m.messages)
}

func (m Model) LastMessage() ChatMessage {
	if len(m.messages) == 0 {
		return ChatMessage{}
	}
	return m.messages[len(m.messages)-1]
}

func (m *Model) HandleEventForTest(ev Event) {
	m.handleEvent(ev)
}

func NewModel(k *kernel.Kernel) Model {
	ta := textarea.New()
	ta.Placeholder = "Type a message and press Enter..."
	ta.SetWidth(80)
	ta.SetHeight(3)
	ta.Focus()
	ta.ShowLineNumbers = false
	ta.KeyMap.InsertNewline.SetEnabled(false)

	evCh := make(chan Event, 256)
	w := NewEventWriter(evCh)
	k.SetStreamWriter(w)

	return Model{
		kernel:    k,
		textarea:  ta,
		events:    evCh,
		writer:    w,
		agentRuns: make(map[string]bool),
	}
}

func (m Model) Init() tea.Cmd {
	return tea.Batch(
		textarea.Blink,
		m.waitEvent(),
		m.tickRefresh(),
		m.flushWriter(),
	)
}

func (m Model) waitEvent() tea.Cmd {
	return func() tea.Msg {
		return <-m.events
	}
}

func (m Model) tickRefresh() tea.Cmd {
	return tea.Tick(time.Second, func(t time.Time) tea.Msg {
		return refreshMsg{}
	})
}

type refreshMsg struct{}
type flushMsg struct{}

func (m Model) flushWriter() tea.Cmd {
	return tea.Tick(150*time.Millisecond, func(t time.Time) tea.Msg {
		m.writer.Flush()
		return flushMsg{}
	})
}

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.updateLayout()

	case tea.KeyMsg:
		switch msg.Type {
		case tea.KeyCtrlC, tea.KeyEsc:
			return m, tea.Quit
		case tea.KeyEnter:
			if m.loading {
				return m, nil
			}
			input := strings.TrimSpace(m.textarea.Value())
			if input == "" {
				return m, nil
			}
			m.textarea.SetValue("")
			m.messages = append(m.messages, ChatMessage{Role: "user", Content: input})
			m.updateViewportContent()
			m.loading = true
			return m, tea.Batch(m.sendMessage(input), m.waitEvent())
		}

	case Event:
		m.handleEvent(msg)
		return m, m.waitEvent()

	case refreshMsg:
		m.refreshAgentStatus()
		return m, m.tickRefresh()

	case flushMsg:
		return m, m.flushWriter()

	case responseMsg:
		m.loading = false
		m.writer.Flush()
		if msg.err != nil {
			m.messages = append(m.messages, ChatMessage{Role: "system", Content: fmt.Sprintf("Error: %v", msg.err)})
			m.updateViewportContent()
		}
		return m, nil
	}

	var cmd tea.Cmd
	m.textarea, cmd = m.textarea.Update(msg)
	cmds = append(cmds, cmd)

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

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

type responseMsg struct {
	content string
	err     error
}

func (m Model) sendMessage(input string) tea.Cmd {
	return func() tea.Msg {
		resp, err := m.kernel.SendMessage("user", "llm", input)
		return responseMsg{content: resp, err: err}
	}
}

func (m *Model) handleEvent(ev Event) {
	switch ev.Type {
	case "stream":
		if len(m.messages) > 0 && m.messages[len(m.messages)-1].Role == "assistant" {
			last := &m.messages[len(m.messages)-1]
			last.Content += ev.Content
		} else {
			m.messages = append(m.messages, ChatMessage{Role: "assistant", Content: ev.Content})
		}
		m.updateViewportContent()
	case "agent_start":
		m.agentRuns[ev.Agent] = true
	case "agent_end":
		m.agentRuns[ev.Agent] = false
	}
}

func (m *Model) refreshAgentStatus() {
	if as := m.kernel.ActorSystem(); as != nil {
		for _, info := range as.AgentInfos() {
			if info.Status == actor.StatusProcessing {
				m.agentRuns[info.ID] = true
			} else {
				m.agentRuns[info.ID] = false
			}
		}
	}
}

func (m *Model) updateLayout() {
	if m.width == 0 || m.height == 0 {
		return
	}

	rightWidth := 42
	leftWidth := m.width - rightWidth - 3
	headerHeight := 3
	mainHeight := m.height - headerHeight - 1
	inputHeight := 5
	chatHeight := mainHeight - inputHeight
	if chatHeight < 10 {
		chatHeight = 10
	}

	m.textarea.SetWidth(leftWidth)
	m.viewport.Width = leftWidth
	m.viewport.Height = chatHeight

	if !m.ready {
		m.viewport = viewport.New(leftWidth, chatHeight)
		m.viewport.SetContent("")
		m.ready = true
	} else {
		m.viewport.Width = leftWidth
		m.viewport.Height = chatHeight
	}
}

func (m *Model) updateViewportContent() {
	var b strings.Builder
	for _, msg := range m.messages {
		b.WriteString(m.formatMessage(msg))
		b.WriteString("\n\n")
	}
	m.viewport.SetContent(b.String())
	m.viewport.GotoBottom()
}

func (m Model) formatMessage(msg ChatMessage) string {
	switch msg.Role {
	case "user":
		return userStyle.Render("You: ") + "\n" + msg.Content
	case "assistant":
		if msg.Agent != "" {
			return agentStyle.Render(fmt.Sprintf("[%s] ", msg.Agent)) + "\n" + msg.Content
		}
		return agentStyle.Render("Assistant: ") + "\n" + msg.Content
	case "system":
		return systemStyle.Render("System: ") + "\n" + msg.Content
	default:
		return msg.Content
	}
}

func (m Model) View() string {
	if m.width == 0 || m.height == 0 {
		return "Initializing..."
	}

	rightWidth := 42
	leftWidth := m.width - rightWidth - 3
	headerHeight := 3
	mainHeight := m.height - headerHeight - 1
	inputHeight := 5
	chatHeight := mainHeight - inputHeight
	if chatHeight < 10 {
		chatHeight = 10
	}

	header := m.renderHeader()

	chatBox := boxStyle.Width(leftWidth).Height(chatHeight).Render(m.viewport.View())
	inputBox := boxStyle.Width(leftWidth).Render(m.textarea.View())
	leftPanel := lipgloss.JoinVertical(lipgloss.Left, chatBox, inputBox)

	stats := m.renderStats()
	agents := m.renderAgents()

	statsHeight := lipgloss.Height(stats)
	agentsHeight := lipgloss.Height(agents)
	remainingHeight := mainHeight - statsHeight - agentsHeight - 2
	if remainingHeight < 3 {
		remainingHeight = 3
	}

	empty := boxStyle.Width(rightWidth - 4).Height(remainingHeight).Render("")
	rightPanel := lipgloss.JoinVertical(lipgloss.Left, stats, agents, empty)

	if m.loading {
		leftPanel = lipgloss.JoinVertical(lipgloss.Left, leftPanel, processingStyle.Render("Processing..."))
	}

	body := lipgloss.JoinHorizontal(lipgloss.Top, leftPanel, rightPanel)
	return lipgloss.JoinVertical(lipgloss.Left, header, body)
}

func (m Model) renderHeader() string {
	version := "v0.1.0"
	title := titleStyle.Render("orca.agent ") + mutedStyle.Render(version)
	return lipgloss.NewStyle().
		Width(m.width).
		Height(1).
		Padding(0, 1).
		Background(lipgloss.Color(colors.bg)).
		Render(title)
}

func (m Model) renderStats() string {
	var b strings.Builder
	b.WriteString(titleStyle.Render("Statistics") + "\n\n")

	tools := 0
	if tm := m.kernel.ToolManager(); tm != nil {
		tools = tm.Count()
	}
	skills := 0
	if sm := m.kernel.SkillManager(); sm != nil {
		skills = len(sm.ListSkills())
	}
	agents := 0
	if as := m.kernel.ActorSystem(); as != nil {
		agents = as.AgentCount()
	}

	b.WriteString(statLabelStyle.Render("Tools: "))
	b.WriteString(statValueStyle.Render(fmt.Sprintf("%d", tools)) + "\n")
	b.WriteString(statLabelStyle.Render("Skills: "))
	b.WriteString(statValueStyle.Render(fmt.Sprintf("%d", skills)) + "\n")
	b.WriteString(statLabelStyle.Render("Agents: "))
	b.WriteString(statValueStyle.Render(fmt.Sprintf("%d", agents)) + "\n")

	return boxStyle.Width(38).Render(b.String())
}

func (m Model) renderAgents() string {
	var b strings.Builder
	b.WriteString(titleStyle.Render("Active Agents") + "\n\n")

	if as := m.kernel.ActorSystem(); as != nil {
		for _, info := range as.AgentInfos() {
			status := "idle"
			style := idleAgentStyle
			if info.Status == actor.StatusProcessing || m.agentRuns[info.ID] {
				status = "running"
				style = activeAgentStyle
			}
			b.WriteString(fmt.Sprintf("• %s: %s\n", info.ID, style.Render(status)))
		}
	}

	return boxStyle.Width(38).Render(b.String())
}