pingo/cmd/pingo.go

// A simple TUI application that charts latency times to specified hosts.
//
// The TUI uses the bubbletea framework for its lifecycle and the lipgloss
// framework for stylized text rendering, for more on these concepts, see their
// respective documentation.
//
// The bubbletea Model consists of a scroll aware container for the pingo.Peak
// bubble, the minimum polling rate to enforce, and a pingo.Peak instance.
//
// The Model.Update function simply interprets ctrl+c signals as quit commands,
// updates viewport/bubble sizes, and runs the viewport and pingo.Peak update
// functions.
//
// The Model.View function prepends and appends a title and footer to the
// viewport, as derived by Model.header() and Model.footer()
//
// The main function first captures the following arguments from the user:
//
//	-s [as time.Duration]: the polling interval
//	-h [as int]: the desired chart height
//	addresses [as []string]: the addresses to poll.
//
// The main function will exit 1 if any of these arguments fail to validate.
//
// Finally, the main function executes bubbletea.Model.Run, exiting 1 upon any
// errors.
package main

import (
	"flag"
	"fmt"
	"os"
	"pingo"
	"time"

	"charm.land/bubbles/v2/viewport"
	tea "charm.land/bubbletea/v2"
	"charm.land/lipgloss/v2"
)

// the bubbletea.Model for the executable
type Model struct {
	viewport viewport.Model
	speed    time.Duration
	p        pingo.Peak
}

// timingMsg is used for tracking Peak.Poll timing intervals, emitted by an
// anonymous function in the update function.
type timingMsg time.Time

// lipgloss styles for the tui
var (
	// footer styles
	titleStyle = lipgloss.NewStyle().
			Align(lipgloss.Center). // implies consumer functions will apply a width
			Italic(true).
			Faint(true)

	// footer style
	footerStyle = lipgloss.NewStyle().
			Align(lipgloss.Center). // implies consumer functions will apply a width
			Italic(true).
			Faint(true)
)

// The tea.Cmd command to be called upon model initialization by bubbletea.
func (m Model) Init() tea.Cmd {
	return tea.Tick(m.speed*time.Millisecond, func(t time.Time) tea.Msg { return timingMsg(t) })
}

// The core Update loop function.
func (m Model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
	// data prep
	var cmds []tea.Cmd // final return
	var cmd tea.Cmd

	switch msg := msg.(type) {
	case tea.WindowSizeMsg: // handle window resizing
		// update pingo.Peak bubble
		m.p.Width = msg.Width
		m.p.Height = msg.Height - m.getVerticalMargin() - 1

		// if the viewport has not been initialized
		if m.viewport.Width() == 0 && m.viewport.Height() == 0 {
			m.viewport = viewport.New( // init viewport
				viewport.WithHeight(msg.Height-m.getVerticalMargin()),
				viewport.WithWidth(msg.Width),
			)
			m.viewport.YPosition = 1 // tell viewport to render after an empty line
		} else {
			m.viewport.SetHeight(msg.Height - m.getVerticalMargin())
			m.viewport.SetWidth(msg.Width)
		}

	case tea.KeyPressMsg: // handle quit
		k := msg.String()
		if k == "ctrl+c" {
			return m, tea.Quit
		}

	case timingMsg: // handle timingMsg (timing loop)
		cmd = tea.Tick(m.speed*time.Millisecond, func(t time.Time) tea.Msg { return timingMsg(t) })
		cmds = append(cmds, tea.Sequence(m.p.Poll(), cmd))
	}

	// update viewport bubble
	m.viewport.SetContent(m.p.View().Content)
	m.viewport, cmd = m.viewport.Update(msg)
	cmds = append(cmds, cmd)

	// update Peak bubble
	m.p, cmd = m.p.Update(msg)
	cmds = append(cmds, cmd)

	return m, tea.Batch(cmds...)
}

func (m Model) View() tea.View {
	m.viewport.SetContent(m.p.View().Content)
	content := fmt.Sprintf("%s%s\n%s", m.header(), m.viewport.View(), m.footer())

	var v tea.View
	v.SetContent(content)
	v.AltScreen = true
	return v
}

func (m Model) header() string { return titleStyle.Width(m.viewport.Width()).Render("pingo v0") }

func (m Model) footer() string {
	return footerStyle.Width(m.viewport.Width()).Render("j/k: down/up\t|\tctrl+c: quit")
}

// get lipgloss.Height value of header and footer
func (m Model) getVerticalMargin() int { return lipgloss.Height(m.header() + m.footer()) }

func main() {
	// get timing interval in milliseconds
	speed := flag.Int("s", 80, "the interval in milliseconds between pings")
	chartHeight := flag.Int("h", 0,
		"the height of the latency chart. set to 0 to render charts full screen.")
	flag.Parse()
	hosts := flag.Args()

	if len(hosts) == 0 {
		fmt.Println("Must specify hosts!")
		return
	}

	if *speed < 1 {
		fmt.Println("speed must not be below 1")
		os.Exit(1)
	}

	p := tea.NewProgram(Model{
		p:     pingo.InitialPeakBubble(hosts, 20, 10, *chartHeight),
		speed: time.Duration(*speed),
	})

	if _, err := p.Run(); err != nil {
		fmt.Printf("Alas, there's been an error: %v", err)
		os.Exit(1)
	}
}