256 lines
5.6 KiB
Go
Raw Normal View History

2025-04-15 19:58:39 +08:00
package statusbar
import (
"fmt"
"lazykimi/pkg/theme"
"strings"
"time"
"github.com/charmbracelet/bubbles/v2/spinner"
tea "github.com/charmbracelet/bubbletea/v2"
"github.com/charmbracelet/lipgloss/v2"
)
type FlashType string
const (
FlashTypeError FlashType = "error"
FlashTypeSuccess FlashType = "success"
)
func FlashSuccess(msg string) tea.Cmd {
return func() tea.Msg {
return EventFlashSuccess(msg)
}
}
func FlashError(msg string) tea.Cmd {
return func() tea.Msg {
return EventFlashError(msg)
}
}
type (
EventFlashSuccess string
EventFlashError string
)
// Model represents the application status bar
type Model struct {
// Layout
width int
styles Style
// Left sections
messagesCount int
modelName string
markdownMode bool
// Right sections
Spinner spinner.Model
isGenerating bool
startTime time.Time
flashType FlashType
flashMessage string
}
// Style represents the styling for the status bar
type Style struct {
base lipgloss.Style
divider lipgloss.Style
message lipgloss.Style
model lipgloss.Style
mode lipgloss.Style
flash map[FlashType]lipgloss.Style
bar lipgloss.Style
}
// New creates a new status bar component
func New(width int) Model {
s := spinner.New()
s.Spinner = spinner.Moon
// Initialize base styles
baseStyle := lipgloss.NewStyle().Bold(true).Padding(0, 1)
return Model{
width: width,
Spinner: s,
styles: Style{
base: baseStyle,
divider: lipgloss.NewStyle().
Foreground(lipgloss.Color("#666666")).
Padding(0, 1),
message: baseStyle.Foreground(theme.Yellow),
model: baseStyle.Foreground(theme.Blue),
mode: baseStyle.Foreground(theme.Green),
flash: map[FlashType]lipgloss.Style{
FlashTypeError: baseStyle.Foreground(theme.Red),
FlashTypeSuccess: baseStyle.Foreground(theme.Green),
},
bar: lipgloss.NewStyle().
BorderStyle(lipgloss.NormalBorder()).
BorderForeground(theme.Gray).
BorderTop(true).
Width(width),
},
}
}
// SetWidth updates the width of the status bar
func (s *Model) SetWidth(width int) {
s.width = width
}
// SetLoading sets loading state
func (s *Model) SetGenerating(loading bool) {
s.isGenerating = loading
if loading {
s.startTime = time.Now()
}
}
func (s *Model) ClearFlashMessage() {
s.flashType = ""
s.flashMessage = ""
}
// SetError sets an error message
func (s *Model) setFlashMessage(t FlashType, msg string) {
s.flashType = t
s.flashMessage = msg
}
// SetMessageCount sets the message count
func (s *Model) SetMessageCount(count int) {
s.messagesCount = count
}
// SetModelName sets the model name
func (s *Model) SetModelName(name string) {
s.modelName = name
}
// SetMarkdownMode sets the markdown mode
func (s *Model) SetMarkdownMode(enable bool) {
s.markdownMode = enable
}
// Init implements tea.Model
func (s *Model) Init() tea.Cmd {
return s.Spinner.Tick
}
// Update implements tea.Model
func (s Model) Update(msg tea.Msg) (Model, tea.Cmd) {
var cmd tea.Cmd
var cmds []tea.Cmd
switch msg := msg.(type) {
case tea.WindowSizeMsg:
s.SetWidth(msg.Width)
case spinner.TickMsg:
s.Spinner, cmd = s.Spinner.Update(msg)
cmds = append(cmds, cmd)
case EventFlashError:
s.setFlashMessage(FlashTypeError, string(msg))
case EventFlashSuccess:
s.setFlashMessage(FlashTypeSuccess, string(msg))
}
return s, tea.Batch(cmds...)
}
// View implements tea.Model
func (s *Model) View() string {
// Render left and right sections
leftBar := s.joinSections(s.getLeftSections())
rightBar := s.joinSections(s.getRightSections())
// Calculate and handle available space
availableWidth := s.width
leftWidth := lipgloss.Width(leftBar)
rightWidth := lipgloss.Width(rightBar)
spacerWidth := availableWidth - leftWidth - rightWidth
var statusBar string
// Handle different space scenarios
switch {
case spacerWidth >= 0:
// Enough space for everything
spacer := strings.Repeat(" ", spacerWidth)
statusBar = fmt.Sprintf("%s%s%s", leftBar, spacer, rightBar)
case availableWidth >= leftWidth:
// Only show left bar
statusBar = leftBar
default:
// Not enough space - truncate left bar
statusBar = lipgloss.NewStyle().
MaxWidth(availableWidth).
Render(leftBar)
}
// Apply the final bar styling
return s.styles.bar.Width(s.width).Render(statusBar)
}
// getLeftSections returns the rendered left sections of the status bar
func (s *Model) getLeftSections() []string {
mode := "PLAIN"
if s.markdownMode {
mode = "MARKDOWN"
}
return []string{
s.styles.message.Render(fmt.Sprintf("💬 %d", s.messagesCount)),
s.styles.model.Render(fmt.Sprintf("🤖 %s", s.modelName)),
s.styles.mode.Render(mode),
}
}
// getRightSections returns the rendered right sections of the status bar
func (s *Model) getRightSections() []string {
var sections []string
// Generation status
if s.isGenerating {
duration := time.Since(s.startTime).Round(time.Millisecond)
sections = append(sections,
s.styles.flash[FlashTypeSuccess].Render(
fmt.Sprintf("⚡ %s %.1fs", s.Spinner.View(), duration.Seconds()),
),
)
}
// Flash message
if s.flashMessage != "" {
var icon string
switch s.flashType {
case FlashTypeError:
icon = "×"
case FlashTypeSuccess:
icon = "✓"
}
sections = append(sections,
s.styles.flash[s.flashType].Render(
fmt.Sprintf("%s %s", icon, s.flashMessage),
),
)
}
return sections
}
// joinSections joins sections with a dot separator
func (s *Model) joinSections(sections []string) string {
if len(sections) == 0 {
return ""
}
separator := s.styles.divider.SetString("┃").String()
result := sections[0]
for _, section := range sections[1:] {
result = lipgloss.JoinHorizontal(lipgloss.Center, result, separator, section)
}
return result
}