llmgw/internal/cli/tui_init.go

package cli

import (
	"context"
	"fmt"
	"io"
	"net/mail"
	"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.
const (
	fAdmin = iota
	fHosts
	fPostgres
	fMirror
	fNode
	fModel
	fGuard
	fStart
	fCount
)

type initModel struct {
	in    *serverInit
	found []string // this machine's addresses

	// The form, until it is submitted.
	form      bool
	focus     int
	admin     textField
	hosts     textField
	toggles   [fCount]bool
	formError string

	// 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 {
	admin := newTextField()
	admin.Placeholder = "you@example.com"
	admin.SetValue(in.admin)
	admin.CharLimit = 200
	admin.Width = 40
	hosts := newTextField()
	hosts.Placeholder = "llm.example.com, 203.0.113.10"
	hosts.SetValue(strings.Join(in.hosts, ", "))
	hosts.CharLimit = 400
	hosts.Width = 40
	sp := newSpinner()
	m := initModel{in: in, found: found, form: askForm, admin: admin, hosts: hosts, spin: sp, width: 100}
	m.toggles[fPostgres] = in.postgres
	m.toggles[fMirror] = !in.noMirror
	m.toggles[fNode] = !in.noNode
	m.toggles[fModel] = !in.noModel
	m.toggles[fGuard] = !in.noGuardModel
	if askForm {
		m.admin.Focus()
	}
	return m
}

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
		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.KeyTab, tea.KeyDown:
		m.moveFocus(1)
		return m, nil
	case tea.KeyShiftTab, tea.KeyUp:
		m.moveFocus(-1)
		return m, nil
	case tea.KeyEnter:
		if m.focus != fStart {
			m.moveFocus(1)
			return m, nil
		}
		admin := strings.TrimSpace(m.admin.Value())
		if admin != "" {
			if _, err := mail.ParseAddress(admin); err != nil {
				m.formError = "the administrator's email is not an address"
				m.focus = fAdmin
				m.admin.Focus()
				return m, nil
			}
		}
		m.in.admin = admin
		m.in.hosts = nil
		for _, h := range strings.Split(m.hosts.Value(), ",") {
			if h = strings.TrimSpace(h); h != "" {
				m.in.hosts = append(m.in.hosts, h)
			}
		}
		m.in.postgres = m.toggles[fPostgres]
		m.in.noMirror = !m.toggles[fMirror]
		m.in.noNode = !m.toggles[fNode]
		m.in.noModel = !m.toggles[fModel]
		m.in.noGuardModel = !m.toggles[fGuard]
		m.form = false
		return m, tea.Batch(m.spin.Tick, func() tea.Msg { return startInitMsg{} })
	}
	switch m.focus {
	case fAdmin:
		var cmd tea.Cmd
		m.admin, cmd = m.admin.Update(k)
		return m, cmd
	case fHosts:
		var cmd tea.Cmd
		m.hosts, cmd = m.hosts.Update(k)
		return m, cmd
	case fPostgres, fMirror, fNode, fModel, fGuard:
		if k.String() == " " || k.String() == "x" {
			m.toggles[m.focus] = !m.toggles[m.focus]
		}
	}
	return m, nil
}

func (m *initModel) moveFocus(d int) {
	m.focus = (m.focus + d + fCount) % fCount
	// With --mirror there is no mirror to run here, so nothing to tick.
	if m.focus == fMirror && m.in.mirrorURL != "" {
		m.focus = (m.focus + d + fCount) % fCount
	}
	m.admin.Blur()
	m.hosts.Blur()
	switch m.focus {
	case fAdmin:
		m.admin.Focus()
	case fHosts:
		m.hosts.Focus()
	}
}

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: "Settings", State: "current"}, {Label: "Install", State: "todo"}}
	keys := "Tab/↑↓ move · Space tick · Enter on Set up · Ctrl+C leave"
	if !m.form {
		steps[1].State, steps[2].State = "done", "current"
		keys = "Ctrl+C stop"
		if m.done {
			steps[2].State = "done"
			keys = "Enter close"
		}
	}
	host, _ := os.Hostname()
	return frame{Width: m.width + navWidth + 6, Height: m.height, Title: "llmgw setup", Where: host,
		NavTitle: "Steps", Nav: steps, PanelTitle: "This machine as a control plane", Body: body, Keys: keys}.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 {
		return b.String() + m.formView()
	}
	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
}

func (m initModel) formView() string {
	var b strings.Builder
	label := func(i int, s string) string {
		if m.focus == i {
			return sTitle.Render("› " + s)
		}
		return "  " + s
	}
	check := func(i int, s, hint string) string {
		box := "[ ]"
		if m.toggles[i] {
			box = "[" + sOK.Render("x") + "]"
		}
		return fmt.Sprintf("%s %s %s\n", label(i, box), s, sDim.Render(hint))
	}
	fmt.Fprintf(&b, "%s\n    %s\n", label(fAdmin, "Administrator's email"), m.admin.View())
	b.WriteString("    " + sDim.Render("A password is generated and shown once. Leave empty to create accounts later.") + "\n\n")
	fmt.Fprintf(&b, "%s\n    %s\n", label(fHosts, "Other names nodes will reach this machine by"), m.hosts.View())
	found := "none found"
	if len(m.found) > 0 {
		found = strings.Join(m.found, ", ")
	}
	b.WriteString("    " + sDim.Render("Already in the certificate: "+found+". Add a DNS name or NAT address here.") + "\n\n")
	b.WriteString(check(fPostgres, "PostgreSQL", "rather than SQLite; installed here, for several control planes or a larger store"))
	if m.in.mirrorURL != "" {
		b.WriteString("  Model mirror " + sDim.Render("at "+m.in.mirrorURL+", from --mirror") + "\n")
	} else {
		b.WriteString(check(fMirror, "Model mirror", "downloads model files once for every node"))
	}
	b.WriteString(check(fNode, "This machine as a node", "serves models here too, on its GPUs or CPU"))
	b.WriteString(check(fModel, "Starter model", "a small Qwen, so there is something to ask at once"))
	b.WriteString(check(fGuard, "Guardrails' judge", "Llama Guard 3 1B on a CPU, refusing violent, criminal and abusive requests"))
	b.WriteString("\n" + label(fStart, lipgloss.NewStyle().Bold(true).Render("Set up")) + "\n")
	if m.formError != "" {
		b.WriteString("\n" + sBad.Render(m.formError) + "\n")
	}
	if !m.framed { // the frame's key bar says it
		b.WriteString("\n" + sDim.Render("tab or ↑↓ to move · space to tick · enter on Set up to start · ctrl+c to leave") + "\n")
	}
	return b.String()
}

// 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
}