package scheduler

import (
	"gitee.com/h79/goutils/common"
	"sync"
)

type Pool struct {
	JobQueue  chan Job
	JobResult chan JobResult
	workers   chan *jobWorker
	stop      chan bool
	wg        sync.WaitGroup
}

func NewPool(numWorkers int, jobQueueLen int) *Pool {
	if numWorkers <= 0 || jobQueueLen <= 0 {
		panic("length is zero")
	}
	pool := &Pool{
		JobQueue:  make(chan Job, jobQueueLen),
		JobResult: make(chan JobResult, jobQueueLen),
		workers:   make(chan *jobWorker, numWorkers),
		stop:      make(chan bool),
	}
	for i := 0; i < cap(pool.workers); i++ {
		worker := newWorker(pool)
		worker.start()
	}
	go pool.dispatch()
	return pool
}

func (p *Pool) dispatch() {
	defer common.Recover()
	for {
		select {
		case job := <-p.JobQueue:
			worker := <-p.workers
			worker.timeout = job.JTimeout
			worker.job <- job
		case stop := <-p.stop:
			if stop {
				for i := 0; i < cap(p.workers); i++ {
					worker := <-p.workers
					worker.stop <- true
					<-worker.stop
				}
				p.stop <- true
				return
			}
		}
	}
}

func (p *Pool) JobDone() {
	p.wg.Done()
}

func (p *Pool) WaitCount(count int) {
	if count != cap(p.JobQueue) {
		count = cap(p.JobQueue)
	}
	p.wg.Add(count)
}

func (p *Pool) WaitAll() {
	p.wg.Wait()
	close(p.JobResult)
}

func (p *Pool) Release() {
	p.stop <- true
	<-p.stop
}