statview/view.go

package statview

import (
	"io"
	"os"
	"sort"
	"sync"
)

type View[K comparable] struct {
	writer   io.Writer
	sortFunc func(i, j K) bool

	lastLineAmount int
	lastReports    map[K]report[K]

	tasks   []K
	taskMap map[K]struct{}

	wg       *sync.WaitGroup
	reportCh chan report[K]
	doneCh   chan struct{}
}

func New[K comparable]() *View[K] {
	return NewForWriter[K](os.Stderr)
}

func NewForWriter[K comparable](w io.Writer) *View[K] {
	return &View[K]{
		writer:  w,
		tasks:   make([]K, 0),
		taskMap: map[K]struct{}{},

		reportCh:    make(chan report[K], 100),
		doneCh:      make(chan struct{}, 1),
		lastReports: map[K]report[K]{},
		wg:          &sync.WaitGroup{},
	}
}

func (v *View[K]) Show() {
	go v.WaitForAllTasks()

	for status := range v.reportCh {
		v.lastReports[status.ID] = status
		v.Add(status.ID)
		v.print()
	}

	v.doneCh <- struct{}{}
}

func (v *View[K]) SortFunc(f func(i, j K) bool) {
	v.sortFunc = f
}

func (v *View[K]) Add(task K) {
	if _, ok := v.taskMap[task]; ok {
		return
	}

	v.wg.Add(1)
	v.tasks = append(v.tasks, task)
	v.taskMap[task] = struct{}{}

	if v.sortFunc != nil {
		sort.Slice(v.tasks, func(i, j int) bool {
			return v.sortFunc(v.tasks[i], v.tasks[j])
		})
	}
}

func (v *View[K]) Report(task K, status string) {
	v.reportCh <- report[K]{ID: task, Text: status}
}

func (v *View[K]) Done(task K) {
	v.wg.Done()
}

func (v *View[K]) WaitForAllTasks() {
	go func() {
		defer close(v.reportCh)
		v.wg.Wait()
	}()

	<-v.doneCh
}