llmgw/internal/cli/tui_node.go

package cli

import (
	"context"
	"fmt"
	"io"
	"os"
	"strings"

	"github.com/charmbracelet/bubbles/spinner"
	"github.com/charmbracelet/bubbles/textinput"
	tea "github.com/charmbracelet/bubbletea"
	"github.com/charmbracelet/lipgloss"
	"github.com/mattn/go-isatty"

	"gitea.mynodes.xyz/ilker/llmgw/internal/onboard"
)

// The full-screen way through 'llmgw node add': the same steps as the plain
// one, drawn as a checklist that fills in, with the questions a script
// cannot ask — which fabric the adapters are on, whether to go on past a
// problem — asked in place.

func isTerminal(w io.Writer) bool {
	f, ok := w.(*os.File)
	return ok && isatty.IsTerminal(f.Fd()) && os.Getenv("TERM") != "dumb"
}

var (
	tuiAccent = lipgloss.AdaptiveColor{Light: "#0B7F7D", Dark: "#3FC2BC"}
	tuiOK     = lipgloss.AdaptiveColor{Light: "#1F8A4C", Dark: "#6FD3A6"}
	tuiWarn   = lipgloss.AdaptiveColor{Light: "#A15C00", Dark: "#E9B35A"}
	tuiBad    = lipgloss.AdaptiveColor{Light: "#B42318", Dark: "#F07C6C"}
	tuiDim    = lipgloss.AdaptiveColor{Light: "#56696C", Dark: "#8FA5A4"}

	sTitle = lipgloss.NewStyle().Bold(true).Foreground(tuiAccent)
	sStep  = lipgloss.NewStyle().Width(14)
	sDim   = lipgloss.NewStyle().Foreground(tuiDim)
	sOK    = lipgloss.NewStyle().Foreground(tuiOK)
	sWarn  = lipgloss.NewStyle().Foreground(tuiWarn)
	sBad   = lipgloss.NewStyle().Foreground(tuiBad)
	sBox   = lipgloss.NewStyle().Border(lipgloss.RoundedBorder()).BorderForeground(tuiAccent).Padding(0, 1)
)

// Messages from the onboarding goroutine to the screen.
type (
	evMsg     onboard.Event
	finishMsg struct {
		node, summary string
		err           error
	}
	askFabricsMsg struct {
		suggested []onboard.Fabric
		reply     chan []onboard.Fabric
	}
	askContinueMsg struct {
		fails int
		reply chan bool
	}
)

type stepState struct {
	level, detail string
}

type nodeModel struct {
	title   string
	steps   map[string]stepState
	notes   []onboard.Event // warnings and failures, with their fixes
	spin    spinner.Model
	cancel  context.CancelFunc
	width   int
	done    bool
	err     error
	node    string
	summary string

	// The question being asked, if any.
	fabrics   *askFabricsMsg
	inputs    []textinput.Model
	focus     int
	confirm   *askContinueMsg
	quitAsked bool
}

func newNodeModel(title string, cancel context.CancelFunc) nodeModel {
	sp := spinner.New()
	sp.Spinner = spinner.MiniDot
	sp.Style = lipgloss.NewStyle().Foreground(tuiAccent)
	return nodeModel{title: title, steps: map[string]stepState{}, spin: sp, cancel: cancel, width: 100}
}

func (m nodeModel) Init() tea.Cmd { return m.spin.Tick }

func (m nodeModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
	switch msg := msg.(type) {
	case tea.WindowSizeMsg:
		m.width = msg.Width
	case spinner.TickMsg:
		var cmd tea.Cmd
		m.spin, cmd = m.spin.Update(msg)
		return m, cmd
	case evMsg:
		e := onboard.Event(msg)
		switch e.Level {
		case onboard.Running, onboard.Done:
			m.steps[e.Step] = stepState{e.Level, e.Detail}
		case onboard.Fail:
			m.steps[e.Step] = stepState{e.Level, e.Detail}
			if e.Step == onboard.StepCheck && strings.Contains(e.Detail, ":") {
				m.notes = append(m.notes, e)
			}
		case onboard.Warn:
			m.notes = append(m.notes, e)
		case onboard.Info:
			if e.Step == onboard.StepEngines {
				m.steps[e.Step] = stepState{onboard.Running, e.Detail}
			}
		}
	case askFabricsMsg:
		m.fabrics = &msg
		m.inputs = nil
		for i, f := range msg.suggested {
			in := textinput.New()
			in.Prompt = ""
			in.SetValue(f.Name)
			in.CharLimit = 40
			in.Width = 24
			if i == 0 {
				in.Focus()
			}
			m.inputs = append(m.inputs, in)
		}
		m.focus = 0
	case askContinueMsg:
		m.confirm = &msg
	case finishMsg:
		m.done, m.err, m.node, m.summary = true, msg.err, msg.node, msg.summary
		return m, tea.Quit
	case tea.KeyMsg:
		return m.key(msg)
	}
	return m, nil
}

func (m nodeModel) key(k tea.KeyMsg) (tea.Model, tea.Cmd) {
	if k.Type == tea.KeyCtrlC {
		m.cancel()
		m.quitAsked = true
		return m, tea.Quit
	}
	switch {
	case m.confirm != nil:
		switch strings.ToLower(k.String()) {
		case "y":
			m.confirm.reply <- true
			m.confirm = nil
		case "n", "enter", "esc":
			m.confirm.reply <- false
			m.confirm = nil
		}
		return m, nil
	case m.fabrics != nil:
		switch k.Type {
		case tea.KeyTab, tea.KeyDown:
			m.move(1)
			return m, nil
		case tea.KeyShiftTab, tea.KeyUp:
			m.move(-1)
			return m, nil
		case tea.KeyEsc:
			m.fabrics.reply <- nil
			m.fabrics = nil
			return m, nil
		case tea.KeyEnter:
			var out []onboard.Fabric
			for i, in := range m.inputs {
				if name := strings.TrimSpace(in.Value()); name != "" {
					out = append(out, onboard.Fabric{Name: name, HCAs: m.fabrics.suggested[i].HCAs})
				}
			}
			m.fabrics.reply <- out
			m.fabrics = nil
			return m, nil
		}
		var cmd tea.Cmd
		m.inputs[m.focus], cmd = m.inputs[m.focus].Update(k)
		return m, cmd
	}
	if k.String() == "q" {
		m.cancel()
		m.quitAsked = true
		return m, tea.Quit
	}
	return m, nil
}

func (m *nodeModel) move(d int) {
	m.inputs[m.focus].Blur()
	m.focus = (m.focus + d + len(m.inputs)) % len(m.inputs)
	m.inputs[m.focus].Focus()
}

func (m nodeModel) View() string {
	var b strings.Builder
	b.WriteString(sTitle.Render(m.title) + "\n\n")
	detailWidth := max(20, m.width-22)
	for _, step := range onboard.Steps {
		st, ok := m.steps[step]
		var mark string
		switch {
		case !ok:
			mark = sDim.Render("·")
		case st.level == onboard.Running && !m.done:
			mark = m.spin.View()
		case st.level == onboard.Fail:
			mark = sBad.Render("✗")
		case st.level == onboard.Running:
			mark = sDim.Render("·")
		default:
			mark = sOK.Render("✓")
		}
		name := sStep.Render(step)
		if !ok {
			name = sDim.Render(sStep.Render(step))
		}
		fmt.Fprintf(&b, "  %s %s %s\n", mark, name, truncate(st.detail, detailWidth))
	}

	if len(m.notes) > 0 {
		b.WriteString("\n")
		for _, n := range m.notes {
			mark := sWarn.Render("!")
			if n.Level == onboard.Fail {
				mark = sBad.Render("✗")
			}
			fmt.Fprintf(&b, "  %s %s\n", mark, wrap(n.Detail, m.width-6, "    "))
			if n.Fix != "" {
				fmt.Fprintf(&b, "    %s\n", sDim.Render("→ "+wrap(n.Fix, m.width-8, "      ")))
			}
		}
	}

	switch {
	case m.fabrics != nil:
		var q strings.Builder
		q.WriteString("These adapters are up. Which fabric is each group on?\n")
		q.WriteString(sDim.Render("Give each the name it has on every other node attached to it.\nLeave a name empty to leave that group out.") + "\n\n")
		for i, f := range m.fabrics.suggested {
			fmt.Fprintf(&q, "  %s  %s\n", m.inputs[i].View(), sDim.Render(strings.Join(f.HCAs, ", ")))
		}
		q.WriteString("\n" + sDim.Render("tab next · enter use these · esc none"))
		b.WriteString("\n" + sBox.Render(q.String()) + "\n")
	case m.confirm != nil:
		q := fmt.Sprintf("doctor found %d problem(s) above.\nAdd the node anyway? %s", m.confirm.fails, sDim.Render("y / N"))
		b.WriteString("\n" + sBox.BorderForeground(tuiWarn).Render(q) + "\n")
	case m.done && m.err != nil:
		b.WriteString("\n" + sBad.Render("✗ ") + wrap(m.err.Error(), m.width-4, "  ") + "\n")
	case m.done:
		b.WriteString("\n" + sOK.Render("✓ ") + m.summary + "\n")
	case m.quitAsked:
		b.WriteString("\n" + sDim.Render("stopped") + "\n")
	default:
		b.WriteString("\n" + sDim.Render("ctrl+c to stop") + "\n")
	}
	return b.String()
}

func truncate(s string, n int) string {
	if n <= 1 || len([]rune(s)) <= n {
		return s
	}
	r := []rune(s)
	return string(r[:n-1]) + "…"
}

// wrap breaks s at spaces to fit width, indenting every line after the first.
func wrap(s string, width int, indent string) string {
	if width < 20 {
		return s
	}
	var lines []string
	line := ""
	for _, w := range strings.Fields(s) {
		if line != "" && len(line)+1+len(w) > width {
			lines = append(lines, line)
			line = w
			continue
		}
		if line == "" {
			line = w
		} else {
			line += " " + w
		}
	}
	lines = append(lines, line)
	return strings.Join(lines, "\n"+indent)
}

// runTUI adds the node with the screen above.
func (na *nodeAdd) runTUI(ctx context.Context) error {
	ctx, cancel := context.WithCancel(ctx)
	defer cancel()
	m := newNodeModel("llmgw · adding "+na.opts.Addr, cancel)
	p := tea.NewProgram(m)

	emit := func(e onboard.Event) { p.Send(evMsg(e)) }
	na.opts.AskFabrics = func(s []onboard.Fabric) []onboard.Fabric {
		reply := make(chan []onboard.Fabric)
		p.Send(askFabricsMsg{s, reply})
		select {
		case f := <-reply:
			return f
		case <-ctx.Done():
			return nil
		}
	}
	na.opts.AskContinue = func(fails int) bool {
		reply := make(chan bool)
		p.Send(askContinueMsg{fails, reply})
		select {
		case ok := <-reply:
			return ok
		case <-ctx.Done():
			return false
		}
	}

	go func() {
		res, err := onboard.New(na.opts, emit).Run(ctx)
		node := ""
		if res != nil {
			node = res.Node
		}
		if err == nil && !na.noWait {
			err = na.watch(ctx, node, emit)
		}
		summary := ""
		if err == nil {
			summary = node + " is in the cluster"
		}
		p.Send(finishMsg{node, summary, err})
	}()

	final, err := p.Run()
	if err != nil {
		return err
	}
	fm := final.(nodeModel)
	if fm.quitAsked && !fm.done {
		return fmt.Errorf("stopped; the node may be half set up, and running this again finishes it")
	}
	return fm.err
}