Update progress bar

* Rename steps to total
* Support concurrent updates
* Add method to watch progress

1 files changed, 84 insertions, 19 deletions

diff --git a/tui/progressbar.go b/tui/progressbar.go
index 0d386d5..7cd0b44 100644
--- a/tui/progressbar.go
+++ b/tui/progressbar.go
@@ -1,38 +1,62 @@
 package tui
 
 import (
+	"context"
 	"fmt"
 	"strings"
+	"sync/atomic"
 	"time"
 )
 
 type ProgressBar struct {
-	title string
-	width int
-	steps int
-	n     int
-	start time.Time
+	title     string
+	width     int
+	total     atomic.Int64
+	completed atomic.Int64
+	start     time.Time
 }
 
-func NewProgressBar(title string, width, steps int, start time.Time) *ProgressBar {
+func NewProgressBar(title string, width, total int, start time.Time) *ProgressBar {
 	pb := &ProgressBar{
 		title: title,
 		width: width,
-		steps: steps,
 		start: start,
 	}
 
+	pb.total.Store(int64(total))
+
 	pb.Print()
 
 	return pb
 }
 
-func (pb *ProgressBar) Increment() {
-	pb.n++
+func (pb *ProgressBar) SetTotal(steps int) {
+	pb.total.Store(int64(steps))
+}
+
+func (pb *ProgressBar) Total() int {
+	return int(pb.total.Load())
+}
+
+func (pb *ProgressBar) SetCompleted(steps int) {
+	pb.completed.Store(int64(steps))
+}
+
+func (pb *ProgressBar) Completed() int {
+	return int(pb.completed.Load())
 }
 
-func (pb *ProgressBar) Update(steps int) {
-	pb.n = steps
+func (pb *ProgressBar) SetProgress(completed, total int) {
+	pb.completed.Store(int64(completed))
+	pb.total.Store(int64(total))
+}
+
+func (pb *ProgressBar) Progress() (int, int) {
+	return pb.Completed(), pb.Total()
+}
+
+func (pb *ProgressBar) Add(steps int) {
+	pb.completed.Add(int64(steps))
 }
 
 func (pb *ProgressBar) Print() {
@@ -44,21 +68,62 @@ func (pb *ProgressBar) Finish() {
 }
 
 func (pb *ProgressBar) String() string {
-	progress := int(float64(pb.n) / float64(pb.steps) * float64(pb.width))
-	percentage := int(float64(pb.n) / float64(pb.steps) * 100)
+	title := pb.title
+	width := pb.width
+	total := pb.Total()
+	completed := pb.Completed()
+	start := pb.start
+
+	var progress int
+	var percentage int
 
-	lap := time.Since(pb.start)
+	if total > 0 {
+		progress = int(float64(completed) / float64(total) * float64(width))
+		percentage = int(float64(completed) / float64(total) * 100)
+	}
+
+	lap := time.Since(start)
 
 	duration := fmt.Sprintf("%02d:%02d:%02d", int(lap.Hours()), int(lap.Minutes())%60, int(lap.Seconds())%60)
 
 	return fmt.Sprintf(
 		"\r[%s] %-20s %3d%% [%s%s] %d/%d",
 		duration,
-		pb.title,
+		title,
 		percentage,
-		strings.Repeat("=", min(progress, pb.width)),
-		strings.Repeat("-", max(0, pb.width-progress)),
-		pb.n,
-		pb.steps,
+		strings.Repeat("=", min(progress, width)),
+		strings.Repeat("-", max(0, width-progress)),
+		completed,
+		total,
 	)
 }
+
+func (pb *ProgressBar) Watch(
+	ctx context.Context,
+	interval time.Duration,
+	callback func() int,
+) {
+	ticker := time.NewTicker(interval)
+
+	defer ticker.Stop()
+
+	for {
+		select {
+		case <-ctx.Done():
+			n := callback()
+
+			pb.SetCompleted(n)
+
+			pb.Finish()
+
+			return
+
+		case <-ticker.C:
+			n := callback()
+
+			pb.SetCompleted(n)
+
+			pb.Print()
+		}
+	}
+}