summaryrefslogtreecommitdiffstats
path: root/progress-bar.go
blob: 11347c3973c9458efdbfbc2457b6f07f2e891840 (plain) (blame)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
package main

import (
	"bufio"
	"errors"
	"flag"
	"fmt"
	"os"
	"strconv"
	"strings"
	"time"

	"github.com/charmbracelet/bubbles/progress"
	tea "github.com/charmbracelet/bubbletea"
	"github.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)
	}
}