package utils

import (
	"sync"
	"sync/atomic"
)

type WaitGroup interface {
	Add(delta int)
	Done()
	Wait()
}

type CountWaitGroup struct {
	wg  sync.WaitGroup
	num atomic.Int64
}

func (c *CountWaitGroup) Add(delta int) {
	c.wg.Add(delta)
	c.num.Add(int64(delta))
}

func (c *CountWaitGroup) Done() {
	c.wg.Done()
	c.num.Add(-1)
}

func (c *CountWaitGroup) Wait() {
	c.wg.Wait()
}

func (c *CountWaitGroup) Count() uint64 {
	return uint64(c.num.Load())
}

func WgWait[T WaitGroup](wg T) <-chan struct{} {
	ch := make(chan struct{})
	go func() {
		wg.Wait()
		ch <- struct{}{}
	}()
	return ch
}