DivaCode/internal/channel/tui.go

package channel

import (
	"context"
	"fmt"
	"log/slog"
	"strings"

	"git.db123.ir/db123/divacode/internal/agent"
	"git.db123.ir/db123/divacode/pkg/types"

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

var (
	userStyle = lipgloss.NewStyle().
			Foreground(lipgloss.Color("#7C3AED")).
			Bold(true)

	assistantStyle = lipgloss.NewStyle().
			Foreground(lipgloss.Color("#10B981")).
			Bold(true)

	reasoningLabelStyle = lipgloss.NewStyle().
				Foreground(lipgloss.Color("#6B7280")).
				Background(lipgloss.Color("#1F2937")).
				Padding(0, 1).
				Italic(true)

	reasoningContentStyle = lipgloss.NewStyle().
				Foreground(lipgloss.Color("#6B7280")).
				Italic(true)

	reasoningHeaderStyle = lipgloss.NewStyle().
				Foreground(lipgloss.Color("#6B7280")).
				Background(lipgloss.Color("#1F2937")).
				Padding(0, 1).
				Bold(true)

	inputStyle = lipgloss.NewStyle().
			BorderStyle(lipgloss.NormalBorder()).
			BorderForeground(lipgloss.Color("#374151"))
)

type sessionState int

const (
	stateReady sessionState = iota
	stateWaiting
)

type model struct {
	agent  *agent.Agent
	ctx    context.Context
	cancel context.CancelFunc

	viewport          viewport.Model
	input             string
	state             sessionState
	messages          []types.Message
	reasoningExpanded map[string]bool
	reasoningOrder    []string
	ready             bool
}

func NewTUIModel(a *agent.Agent) tea.Model {
	ctx, cancel := context.WithCancel(context.Background())
	return &model{
		agent:             a,
		ctx:               ctx,
		cancel:            cancel,
		state:             stateReady,
		reasoningExpanded: make(map[string]bool),
	}
}

func (m *model) Init() tea.Cmd {
	return nil
}

func (m *model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
	switch msg := msg.(type) {
	case tea.WindowSizeMsg:
		if !m.ready {
			m.viewport = viewport.New(msg.Width, msg.Height-3)
			m.viewport.YPosition = 0
			m.ready = true
		} else {
			m.viewport.Width = msg.Width
			m.viewport.Height = msg.Height - 3
		}
		return m, nil

	case tea.MouseMsg:
		var cmd tea.Cmd
		m.viewport, cmd = m.viewport.Update(msg)
		return m, cmd

	case tea.KeyMsg:
		switch msg.String() {
		case "ctrl+c", "ctrl+q":
			m.cancel()
			return m, tea.Quit

		case "enter":
			if m.state == stateReady && strings.TrimSpace(m.input) != "" {
				input := strings.TrimSpace(m.input)
				m.input = ""
				m.state = stateWaiting
				return m, m.sendMessage(input)
			}

		case "tab":
			for i := len(m.messages) - 1; i >= 0; i-- {
				if m.messages[i].Role == types.RoleAssistant && m.messages[i].Reasoning != "" {
					id := m.messages[i].ID
					m.reasoningExpanded[id] = !m.reasoningExpanded[id]
					break
				}
			}

		case "backspace":
			if len(m.input) > 0 {
				m.input = m.input[:len(m.input)-1]
			}

		default:
			if m.state == stateReady {
				m.input += msg.String()
			}
		}

	case responseMsg:
		m.state = stateReady
		if msg.err != nil {
			m.messages = append(m.messages, types.Message{
				Role:    types.RoleSystem,
				Content: fmt.Sprintf("error: %v", msg.err),
			})
		} else {
			newMsgs := m.agent.Messages()
			for _, msg := range newMsgs {
				if msg.Role == types.RoleAssistant && msg.Reasoning != "" {
					found := false
					for _, id := range m.reasoningOrder {
						if id == msg.ID {
							found = true
							break
						}
					}
					if !found {
						m.reasoningOrder = append(m.reasoningOrder, msg.ID)
					}
				}
			}
			m.messages = newMsgs
		}
		m.viewport.GotoBottom()
	}

	return m, nil
}

func (m *model) View() string {
	if !m.ready {
		return "initializing divacode..."
	}

	var b strings.Builder
	for _, msg := range m.messages {
		prefix := "sys"
		style := lipgloss.NewStyle()
		switch msg.Role {
		case types.RoleUser:
			prefix = "you"
			style = userStyle
		case types.RoleAssistant:
			prefix = "diva"
			style = assistantStyle
		default:
			prefix = "system"
		}
		b.WriteString(style.Render(prefix+": ") + msg.Content + "\n")

		if msg.Role == types.RoleAssistant && msg.Reasoning != "" {
			if m.reasoningExpanded[msg.ID] {
				b.WriteString(reasoningHeaderStyle.Render(" reasoning ") + "\n")
				b.WriteString(reasoningContentStyle.Render(msg.Reasoning) + "\n")
			} else {
				label := reasoningLabelStyle.Render(" ` reasoning (tab)")
				b.WriteString(label + "\n")
			}
		}

		b.WriteString("\n")
	}

	m.viewport.SetContent(b.String())

	prompt := "> "
	if m.state == stateWaiting {
		prompt = "... "
	}

	return fmt.Sprintf("%s\n%s%s", m.viewport.View(), inputStyle.Render(prompt), m.input)
}

type responseMsg struct {
	err error
}

func (m *model) sendMessage(content string) tea.Cmd {
	return func() tea.Msg {
		slog.Info("tui message", "content", content)
		_, err := m.agent.HandleMessage(m.ctx, content)
		return responseMsg{err: err}
	}
}