Bubbletea to make it pretty

1 files changed, 129 insertions, 0 deletions

diff --git a/progress-bar.go b/progress-bar.go
new file mode 100644
index 0000000..82e7cb0
--- /dev/null
+++ b/progress-bar.go
@@ -0,0 +1,129 @@
+package main
+
+import (
+	"bufio"
+	"errors"
+	"flag"
+	"fmt"
+	"os"
+	"strconv"
+	"strings"
+	"time"
+
+	"github.com/charmbracelet/bubbles/progress"
+	tea "github.com/charmbracelet/bubbletea"
+	"gitlab.com/rhogenson/deque"
+)
+
+const (
+	padding  = 2
+	maxWidth = 80
+)
+
+type measurement struct {
+	t time.Time
+	i int
+}
+
+var (
+	maxVal       int
+	bar          progress.Model
+	measurements deque.Deque[measurement]
+	eta          time.Duration = -1
+)
+
+type (
+	measurementMsg int
+	quitMsg        struct{}
+)
+
+type model struct{}
+
+func (model) Init() tea.Cmd {
+	return nil
+}
+
+func (m model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
+	switch msg := msg.(type) {
+	case tea.WindowSizeMsg:
+		bar.Width = msg.Width - padding*2
+		return m, nil
+
+	case measurementMsg:
+		n := int(msg)
+		if measurements.Len() > 0 && measurements.At(measurements.Len()-1).i > n {
+			measurements.Reset()
+			eta = -1
+		}
+		now := time.Now()
+		for measurements.Len() > 2 && now.Sub(measurements.At(0).t) > 2*time.Minute {
+			measurements.PopFront()
+		}
+		measurements.PushBack(measurement{now, n})
+		first := measurements.At(0)
+		deltaT := now.Sub(first.t)
+		delta := n - first.i
+		if delta != 0 {
+			eta = time.Duration(float64(maxVal-n) / float64(delta) * float64(deltaT))
+		}
+		return m, bar.SetPercent(float64(n) / float64(maxVal))
+
+	case quitMsg:
+		return m, tea.Quit
+
+	// FrameMsg is sent when the progress bar wants to animate itself
+	case progress.FrameMsg:
+		progressModel, cmd := bar.Update(msg)
+		bar = progressModel.(progress.Model)
+		return m, cmd
+
+	default:
+		return m, nil
+	}
+}
+
+func (m model) View() string {
+	etaStr := "calculating..."
+	if eta >= 0 {
+		etaStr = eta.Round(time.Second).String()
+	}
+	pad := strings.Repeat(" ", padding)
+	return "\n" +
+		pad + bar.View() + "\n" +
+		pad + "ETA: " + etaStr + "\n"
+}
+
+func run() error {
+	flag.Parse()
+	args := flag.Args()
+	if len(args) != 1 {
+		return errors.New("usage error")
+	}
+	var err error
+	maxVal, err = strconv.Atoi(args[0])
+	if err != nil {
+		return err
+	}
+	bar = progress.New(progress.WithDefaultGradient())
+	program := tea.NewProgram(model{}, tea.WithInput(nil))
+	go func() {
+		scanner := bufio.NewScanner(os.Stdin)
+		for scanner.Scan() {
+			n, err := strconv.Atoi(scanner.Text())
+			if err != nil {
+				continue
+			}
+			program.Send(measurementMsg(n))
+		}
+		program.Send(quitMsg{})
+	}()
+	_, err = program.Run()
+	return err
+}
+
+func main() {
+	if err := run(); err != nil {
+		fmt.Fprintln(os.Stderr, err)
+		os.Exit(1)
+	}
+}