util.go

package main

import (
	"bytes"
	"context"
	"fmt"
	"io"
	"os"
	"log/slog"
	"sync"
	"time"
)

func CountLines(path string) int64 {
	fp, err := os.Open(path)
	if err != nil {
		slog.Error("Failed to open file", "path", path, "err", err)
		return 0
	}
	defer fp.Close()

	buf := make([]byte, 32 * 1024)
	newline := []byte{'\n'}
	res := int64(0)

	for {
		size, err := fp.Read(buf)
		res += int64(bytes.Count(buf[:size], newline))

		switch {
		case err == io.EOF:
			return res

		case err != nil:
			slog.Error("Failed to read file", "path", path, "err", err)
			return 0
		}
	}
}

func Drain[T any](ch <-chan T, timeout time.Duration) {
	ctx, cancel := context.WithTimeout(context.Background(), timeout)
	defer cancel()
	for {
		select {
		case <-ctx.Done():
			return
		case _, ok := <-ch:
			if !ok {
				return
			}
		}
	}
}

func ShowProgress[T any](
	progress int64,
	count int64,
	in <-chan T,
	out chan<- T,
	wg *sync.WaitGroup,
) {
	var value T
	prevProgress := progress
	prevTime := time.Now()
	line := "\033[2K\r%v | %d/%d | %d it/s | %s"
	for value = range in {
		progress++
		if progress % 100 == 0 {
			t := time.Now()
			elapsed := t.Sub(prevTime).Seconds()
			if elapsed > 1.0 {
				speed := float64(progress - prevProgress) / elapsed
				timeRemaining := time.Second * time.Duration(
					int64(float64(count - progress) / speed),
				)
				fmt.Printf(
					line,
					value, progress, count, int(speed), timeRemaining.String(),
				)
				prevProgress = progress
				prevTime = t
			}
		}
		out <- value
	}
	elapsed := time.Now().Sub(prevTime).Seconds()
	speed := int(float64(progress - prevProgress) / elapsed)
	fmt.Printf(line, value, progress, count, speed, "done\n")
	close(out)
	wg.Done()
}