package syncz

import (
	"sync"
	"time"
)

type WaitGroup struct {
	mu   sync.Mutex
	cond *sync.Cond
	n    uint32
}

func NewWaitGroup() *WaitGroup {
	wg := &WaitGroup{}
	wg.cond = sync.NewCond(&wg.mu)
	return wg
}

func (wg *WaitGroup) Add() {
	wg.mu.Lock()
	wg.n++
	wg.mu.Unlock()
}

func (wg *WaitGroup) Done() {
	wg.mu.Lock()
	defer wg.mu.Unlock()
	if wg.n == 0 {
		panic("Done() called without calling Add()")
	}
	wg.n--
	if wg.n == 0 {
		wg.cond.Broadcast() // Signal all waiters
	}
}

// Wait waits until the internal counter reaches zero or the timeout occurs.
// Returns false on timeout and true otherwise.
// Returns false immediately if timeout is non-positive.
func (wg *WaitGroup) Wait(timeout time.Duration) bool {
	if timeout <= 0 {
		return false
	}
	timeoutElapsed := make(chan struct{})
	t := time.AfterFunc(timeout, func() {
		close(timeoutElapsed)
		wg.cond.Broadcast() // Signal the loop below
	})
	defer t.Stop()

	wg.mu.Lock()
	defer wg.mu.Unlock()
	for {
		if wg.n == 0 {
			return true
		}
		select {
		case <-timeoutElapsed:
			return false
		default:
		}
		wg.cond.Wait()
	}
}
