llm-shell-hint/tui.go

package main

import (
	"fmt"
	"os"
	"strings"

	"github.com/charmbracelet/bubbles/list"
	"github.com/charmbracelet/bubbles/spinner"
	tea "github.com/charmbracelet/bubbletea"
	"github.com/charmbracelet/lipgloss"
	"golang.org/x/term"
)

type model struct {
	config            *Config
	query             string
	verbose           bool
	commands          []CommandWithComment
	loading           bool
	status            string
	err               error
	spinner           spinner.Model
	list              list.Model
	selectionConfirmed bool
}

func initialModel(config *Config, query string, verbose bool) model {
	sp := spinner.New()
	sp.Spinner = spinner.Dot
	sp.Style = lipgloss.NewStyle().Foreground(lipgloss.Color("205"))

	// Initialize list with empty items
	delegate := NewWrappingDelegate()
	// Set colors based on scheme
	var (
		textColor       = "16"    // Black for light, white for dark
		accentColor     = "33"    // Bright blue
		secondaryColor  = "240"   // Dark gray
	)
	
	if config.ColorScheme == "dark" {
		textColor = "15" // White
	}

	delegate.Styles.SelectedTitle = lipgloss.NewStyle().
		Foreground(lipgloss.Color(textColor)).
		Bold(true).
		BorderLeft(true).
		BorderStyle(lipgloss.NormalBorder()).
		BorderForeground(lipgloss.Color(accentColor)).
		PaddingLeft(1)
	delegate.Styles.SelectedDesc = lipgloss.NewStyle().
		Foreground(lipgloss.Color(secondaryColor)).
		BorderLeft(true).
		BorderStyle(lipgloss.NormalBorder()).
		BorderForeground(lipgloss.Color(accentColor)).
		PaddingLeft(1)
	
	// Customize the delegate to wrap both commands and comments
	delegate.Styles.NormalTitle = lipgloss.NewStyle().
		PaddingLeft(2) // Extra space for inactive items
	
	// Use full terminal width with some padding
	width := 80 // Default width if we can't get terminal size
	if w, _, err := term.GetSize(int(os.Stdout.Fd())); err == nil && w > 0 {
		width = w - 4 // Subtract padding
	}

	l := list.New([]list.Item{}, delegate, width, 20)
	l.SetShowTitle(false)
	l.SetShowStatusBar(false)
	l.SetFilteringEnabled(false)
	l.Title = "" // Set empty title
	l.Styles.PaginationStyle = lipgloss.NewStyle().
		MarginLeft(4).
		Foreground(lipgloss.Color("240")) // Dark gray text
	l.Styles.HelpStyle = lipgloss.NewStyle().
		MarginLeft(4).
		Foreground(lipgloss.Color("240")) // Dark gray text
	
	// Create initial model
	m := model{
		config:  config,
		query:   query,
		verbose: verbose,
		loading: true,
		status:  "Initializing...",
		spinner: sp,
		list:    l,
	}

	// Calculate maximum lines needed for commands
	maxLines := 1
	for _, cmd := range m.commands {
		lines := strings.Count(cmd.Command, "\n") + 1
		if lines > maxLines {
			maxLines = lines
		}
	}
	// Set height based on maximum lines, with a minimum of 3
	if maxLines < 3 {
		maxLines = 3
	}
	
	// Update delegate height
	delegate.SetHeight(maxLines)
	m.list.SetDelegate(delegate)

	return m
}

func (m model) Init() tea.Cmd {
	return tea.Batch(
		m.spinner.Tick,
		func() tea.Msg {
			m.status = "Building prompt..."
			prompt := buildPrompt(m.config, m.query)
			
			m.status = "Querying LLM..."
			commands, err := queryLLM(m.config, prompt, m.verbose)
			if err != nil {
				return err
			}
			
			return commands
		},
	)
}

func (m model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
	switch msg := msg.(type) {
	case tea.KeyMsg:
		switch msg.String() {
		case "ctrl+c", "q":
			return m, tea.Quit
		case "enter":
			if _, ok := m.list.SelectedItem().(item); ok {
				m.selectionConfirmed = true
				return m, tea.Quit
			}
		}
		// Let the list handle its own key messages
		var cmd tea.Cmd
		m.list, cmd = m.list.Update(msg)
		return m, cmd
	case []CommandWithComment:
		m.loading = false
		m.commands = msg
		// Convert commands to list items
		items := make([]list.Item, len(msg))
		for i, cmd := range msg {
			items[i] = item{
				command: cmd.Command,
				comment: cmd.Comment,
			}
		}
		m.list.SetItems(items)
	case error:
		m.err = msg
		return m, tea.Quit
	case spinner.TickMsg:
		var cmd tea.Cmd
		m.spinner, cmd = m.spinner.Update(msg)
		return m, cmd
	}
	return m, nil
}

// Custom list item type
type item struct {
	command string
	comment string
}

func (i item) Title() string       { return i.command }
func (i item) Description() string { return i.comment }
func (i item) FilterValue() string { return i.command }

func (m model) View() string {
	if m.err != nil {
		return fmt.Sprintf("Error: %v\nPress q to quit", m.err)
	}
	
	if m.loading {
		return fmt.Sprintf("\n %s %s...\n", m.spinner.View(), m.status)
	}
	
	style := lipgloss.NewStyle().
		Margin(1, 2)
	
	return style.Render(m.list.View())
}

func runTUI(config *Config, query string, verbose bool) (string, error) {
	// Initialize and run the TUI
	tty, err := os.OpenFile("/dev/tty", os.O_WRONLY, 0)
	if err != nil {
		return "", err
	}
	defer tty.Close()
	p := tea.NewProgram(initialModel(config, query, verbose), tea.WithOutput(tty))
	
	// Run the program and get the final model
	finalModel, err := p.Run()
	if err != nil {
		return "", fmt.Errorf("error running program: %v", err)
	}
	
	// Only print the selected command if we exited with Enter
	if m, ok := finalModel.(model); ok {
		if selected, ok := m.list.SelectedItem().(item); ok {
			if m.selectionConfirmed {
				return selected.command, nil
			}
		}
	}
	return "", nil
}