package utils

import "fmt"

// GoroutinePool define routine pool
type GoroutinePool struct {
	Queue  chan func() error
	Number int
	Total  int

	result         chan error
	finishCallback func()
}

// Init routine pool
func (g *GoroutinePool) Init(number int, total int) {
	g.Queue = make(chan func() error, total)
	g.Number = number
	g.Total = total
	g.result = make(chan error, total)
}

// Start routine pool
func (g *GoroutinePool) Start() {
	for i := 0; i < g.Number; i++ {
		go func() {
			for {
				task, ok := <-g.Queue
				if !ok {
					break
				}

				err := task()
				g.result <- err
			}
		}()
	}

	// get each work result
	for j := 0; j < g.Total; j++ {
		res, ok := <-g.result
		if !ok {
			break
		}

		if res != nil {
			fmt.Println(res)
		}
	}

	if g.finishCallback != nil {
		g.finishCallback()
	}
}

// Stop close queue
func (g *GoroutinePool) Stop() {
	close(g.Queue)
	close(g.result)
}

// AddTask job func
func (g *GoroutinePool) AddTask(task func() error) {
	g.Queue <- task
}

// SetFinishCallback func
func (g *GoroutinePool) SetFinishCallback(callback func()) {
	g.finishCallback = callback
}