package util

import (
	"runtime"
	"sync"
)

type WaitGroup struct {
	q  chan struct{}
	wg sync.WaitGroup
}

// NewWaitGroup 生成
func NewWaitGroup(size int) *WaitGroup {
	if size <= 0 {
		size = runtime.NumCPU()
	}
	return &WaitGroup{
		q: make(chan struct{}, size),
	}
}

// Add 添加或减少任务
func (p *WaitGroup) Add(delta int) {
	var i int
	if delta >= 0 {
		for i = 0; i < delta; i++ {
			p.q <- struct{}{}
			p.wg.Add(1)
		}
	} else {
		for i = 0; i > delta; i-- {
			select {
			case <-p.q:
				p.wg.Add(-1)
			default:
				i = delta
			}
		}
	}
}

// Done 完成任务
func (p *WaitGroup) Done() {
	p.wg.Done()
	select {
	case <-p.q:
	default:
	}
}

// Wait 等待任务完成
func (p *WaitGroup) Wait() {
	p.wg.Wait()
	for cleared := false; !cleared; {
		select {
		case <-p.q:
		default:
			cleared = true
		}
	}
}