llmgw/internal/cli/tui_init.go

package cli

import (
	"context"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strings"
	"sync"

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

// The full-screen way through 'llmgw server init': a short form for what
// init would otherwise take as flags, then init's own lines as a list that
// fills in, with the administrator's password kept in view — and printed
// again once the screen closes, so it is in the terminal's history.

// initLineMsg is one line init wrote.
type initLineMsg string

// initDoneMsg is init finishing.
type initDoneMsg struct{ err error }

// lineWriter turns init's output into one message per line.
type lineWriter struct {
	mu   sync.Mutex
	buf  []byte
	send func(string)
}

func (w *lineWriter) Write(p []byte) (int, error) {
	w.mu.Lock()
	defer w.mu.Unlock()
	w.buf = append(w.buf, p...)
	for {
		i := strings.IndexByte(string(w.buf), '\n')
		if i < 0 {
			return len(p), nil
		}
		w.send(string(w.buf[:i]))
		w.buf = w.buf[i+1:]
	}
}

// Form fields, in order.
type initModel struct {
	in    *serverInit
	found []string // this machine's addresses

	// The parts to install, until they are chosen and the plan agreed to.
	form    bool
	plan    bool
	choices *installForm
	// scroll is how far the panel is scrolled, back to the top on each page.
	scroll int

	// Then the run.
	running  bool
	lines    []string // init's lines, as it wrote them
	password string
	next     []string // the closing "Next:" text
	spin     spinnerModel
	width    int
	height   int
	// framed draws it in the full-screen frame; off when another screen
	// shows it in its own panel.
	framed bool
	done   bool
	err    error
	quit   bool
	start  func()
}

func newInitModel(in *serverInit, found []string, askForm bool) initModel {
	host, _ := os.Hostname()
	return initModel{in: in, found: found, form: askForm, choices: newInstallForm(in, host, found, false),
		spin: newSpinner(), width: 100, height: 30}
}

func (m initModel) Init() tea.Cmd {
	if m.form {
		return nil
	}
	return tea.Batch(m.spin.Tick, func() tea.Msg { return startInitMsg{} })
}

type startInitMsg struct{}

func (m initModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
	switch msg := msg.(type) {
	case tea.WindowSizeMsg:
		m.width, m.height = msg.Width, msg.Height
		m.choices.height = msg.Height
		if m.framed {
			m.width = msg.Width - navWidth - 6
		}
	case spinTickMsg:
		var cmd tea.Cmd
		m.spin, cmd = m.spin.Update(msg)
		return m, cmd
	case startInitMsg:
		m.running = true
		if m.start != nil {
			m.start()
		}
	case initLineMsg:
		m.addLine(string(msg))
	case initDoneMsg:
		m.done, m.err = true, msg.err
		if m.framed {
			return m, nil // left on screen until a key: it is gone once closed
		}
		return m, tea.Quit
	case tea.KeyMsg:
		if msg.Type == tea.KeyCtrlC {
			m.quit = true
			return m, tea.Quit
		}
		if m.done && m.framed && (msg.Type == tea.KeyEnter || msg.String() == "q") {
			return m, tea.Quit
		}
		if m.form {
			return m.formKey(msg)
		}
	}
	return m, nil
}

// addLine files one of init's lines: the password aside, the closing text
// aside, a progress line replacing the one before it.
func (m *initModel) addLine(l string) {
	t := strings.TrimSpace(l)
	switch {
	case t == "":
		return
	case strings.HasPrefix(t, "password: "):
		m.password = strings.TrimPrefix(t, "password: ")
		return
	case strings.HasPrefix(t, "shown once"):
		return
	case !strings.HasPrefix(l, "  ") || len(m.next) > 0:
		// Unindented lines are the closing text, and everything after its
		// first line belongs to it.
		m.next = append(m.next, l)
		return
	}
	if strings.HasPrefix(t, "..") && len(m.lines) > 0 && strings.HasPrefix(strings.TrimSpace(m.lines[len(m.lines)-1]), "..") {
		m.lines[len(m.lines)-1] = l
		return
	}
	m.lines = append(m.lines, l)
}

func (m initModel) formKey(k tea.KeyMsg) (tea.Model, tea.Cmd) {
	switch k.Type {
	case tea.KeyPgDown:
		m.scroll += max(1, m.height-8)
		return m, nil
	case tea.KeyPgUp:
		m.scroll = max(0, m.scroll-max(1, m.height-8))
		return m, nil
	}
	if m.plan {
		switch k.Type {
		case tea.KeyEsc:
			m.plan, m.scroll = false, 0
		case tea.KeyEnter:
			m.form, m.plan, m.scroll = false, false, 0
			return m, tea.Batch(m.spin.Tick, func() tea.Msg { return startInitMsg{} })
		}
		return m, nil
	}
	if m.choices.key(k, m.in) {
		m.plan, m.scroll = true, 0
	}
	return m, nil
}

// command is this install as a command line, for doing it again from a script.
func (in *serverInit) command() string {
	return strings.Join(append([]string{"sudo llmgw server init"}, in.initFlags()...), " \\\n    ")
}

func (m initModel) View() string {
	if !m.framed {
		return m.plainView()
	}
	body := m.plainView()
	if i := strings.Index(body, "\n\n"); i >= 0 {
		body = body[i+1:]
	}
	steps := []navItem{{Label: "Start", State: "done"}, {Label: "Choices", State: "current"},
		{Label: "Plan", State: "todo"}, {Label: "Install", State: "todo"}}
	keys := m.choices.keys()
	title := "What to install"
	switch {
	case m.form && m.plan:
		steps[1].State, steps[2].State = "done", "current"
		keys, title = "Enter install · Esc back to the choices · PgUp/PgDn scroll · Ctrl+C leave", "Plan"
	case !m.form:
		steps[1].State, steps[2].State, steps[3].State = "done", "done", "current"
		keys, title = "Ctrl+C stop", "Installing"
		if m.done {
			steps[3].State = "done"
			keys = "Enter close"
		}
	}
	f := frame{Width: m.width + navWidth + 6, Height: m.height, Title: "llmgw setup",
		NavTitle: "Steps", Nav: steps, PanelTitle: title, Body: body, Keys: keys, Scroll: m.scroll}
	if m.form && !m.plan {
		f.DialogTitle, f.Dialog = m.choices.overlay()
	}
	f.Where, _ = os.Hostname()
	return f.render()
}

// plainView is the screen as text, for a panel or a terminal that is not one.
func (m initModel) plainView() string {
	var b strings.Builder
	b.WriteString(sTitle.Render("llmgw · setting up this machine as a control plane") + "\n\n")
	if m.form && m.plan {
		b.WriteString(m.choices.planView(m.width))
		b.WriteString("\n  " + sDim.Render("The same from a script:") + "\n")
		for _, l := range strings.Split(m.in.command(), "\n") {
			b.WriteString("  " + l + "\n")
		}
		b.WriteString("\n" + button(true, "Enter to install") + "\n")
		return b.String()
	}
	if m.form {
		b.WriteString(m.choices.view(m.width))
		if title, d := m.choices.overlay(); d != "" && !m.framed {
			b.WriteString("\n" + sTitle.Render(title) + "\n" + d + "\n")
		}
		return b.String()
	}
	for _, l := range m.visibleLines() {
		b.WriteString(m.lineView(l) + "\n")
	}
	if m.password != "" {
		box := sBox.BorderForeground(tuiWarn).Render(
			fmt.Sprintf("Administrator %s\npassword  %s\n%s", m.in.admin,
				lipgloss.NewStyle().Bold(true).Render(m.password),
				sDim.Render("Shown once. Copy it now; change it on the account page after signing in.")))
		b.WriteString("\n" + box + "\n")
	}
	switch {
	case m.done && m.err != nil:
		b.WriteString("\n" + sBad.Render("✗ ") + wrap(m.err.Error(), m.width-4, "  ") + "\n" +
			sDim.Render("  It is safe to run init again once that is fixed: it finishes what is missing.") + "\n")
	case m.done:
		b.WriteString("\n" + sOK.Render("✓ ") + "The control plane is up. Add a node: " +
			lipgloss.NewStyle().Bold(true).Render("sudo llmgw node add <user>@<address>") + "\n")
		b.WriteString("  " + sDim.Render("Back up "+m.in.pkiDir+": its ca-key.pem signs every node's identity, and nothing recovers it.") + "\n")
	case m.quit:
		b.WriteString("\n" + sDim.Render("stopped; run init again to finish") + "\n")
	case m.running:
		b.WriteString("\n" + sDim.Render("ctrl+c to stop") + "\n")
	}
	return b.String()
}

// visibleLines are the last lines, enough to follow without the screen
// outgrowing the terminal: the whole log is printed when it closes.
func (m initModel) visibleLines() []string {
	const keep = 16
	if len(m.lines) > keep {
		return m.lines[len(m.lines)-keep:]
	}
	return m.lines
}

func (m initModel) lineView(l string) string {
	t := strings.TrimSpace(l)
	word, rest, _ := strings.Cut(t, " ")
	rest = truncate(strings.TrimSpace(rest), max(20, m.width-6))
	switch word {
	case "ok":
		return "  " + sOK.Render("✓") + " " + rest
	case "kept":
		return "  " + sDim.Render("· "+rest)
	case "skip":
		return "  " + sDim.Render("- "+rest)
	case "!!":
		return "  " + sWarn.Render("! ") + rest
	case "..":
		if m.done {
			return "  " + sDim.Render("· "+rest)
		}
		return "  " + m.spin.View() + " " + rest
	}
	return "    " + t
}

// runInitTUI runs init with the screen above, asking the form first when
// the flags left the administrator unset.
func (in *serverInit) runInitTUI(ctx context.Context, out io.Writer) error {
	var found []string
	for _, ip := range localAddrs() {
		found = append(found, ip.String())
	}
	m := newInitModel(in, found, in.admin == "")
	m.framed = true
	var p *tea.Program

	var (
		logMu sync.Mutex
		log   []string
	)
	w := &lineWriter{send: func(l string) {
		logMu.Lock()
		log = append(log, l)
		logMu.Unlock()
		p.Send(initLineMsg(l))
	}}
	ctx, cancel := context.WithCancel(ctx)
	defer cancel()
	started := make(chan struct{})
	m.start = func() {
		close(started)
		go func() {
			in.out = w
			err := in.run(ctx)
			p.Send(initDoneMsg{err})
		}()
	}
	p = tea.NewProgram(m, tea.WithOutput(out), tea.WithAltScreen())
	final, err := p.Run()
	if err != nil {
		return err
	}
	fm := final.(initModel)
	select {
	case <-started:
	default:
		return fmt.Errorf("left before anything was done")
	}
	if fm.quit && !fm.done {
		cancel()
		return fmt.Errorf("stopped; run init again to finish what was started")
	}

	// The whole run is kept, root's only, with the password taken out: it is
	// shown once, on the screen, and a file on disk would outlive that. On a
	// failure the run is printed too, for the lines the screen had no room
	// for; otherwise the screen's last frame says it, and only what must
	// outlive the screen is printed again: the password, and what to do next.
	logMu.Lock()
	defer logMu.Unlock()
	whole := strings.Join(log, "\n") + "\n"
	kept := whole
	if fm.password != "" {
		kept = strings.ReplaceAll(kept, fm.password, "[password shown once, not kept]")
	}
	saved := ""
	if in.stateDir != "" {
		p := filepath.Join(in.stateDir, "init.log")
		if err := os.WriteFile(p, []byte(kept), 0o600); err == nil {
			saved = p
		}
	}
	fmt.Fprintln(out)
	if fm.err != nil {
		fmt.Fprint(out, whole)
		return fm.err
	}
	if fm.password != "" {
		fmt.Fprintf(out, "Administrator %s, password %s (shown once)\n\n", in.admin, fm.password)
	}
	for _, l := range fm.next {
		fmt.Fprintln(out, l)
	}
	if saved != "" {
		fmt.Fprintf(out, "\nThe whole run is in %s.\n", saved)
	}
	return fm.err
}