package grouptask

import (
	"errors"
	"sync"
)

var (
	ErrTaskPoolIsFull = errors.New("task pool is full")
)

type TaskHandle func(interface{}, uint)

type Tasks struct {
	tasknum  uint
	handler  TaskHandle
	taskpool chan interface{}
	wg       sync.WaitGroup
}

func NewTasks(routine uint, handle TaskHandle) *Tasks {
	return &Tasks{
		tasknum:  routine,
		handler:  handle,
		taskpool: make(chan interface{}, 1000000),
	}
}

func (t *Tasks) AddTask(task interface{}) error {
	select {
	case t.taskpool <- task:
		return nil
	default:
		return ErrTaskPoolIsFull
	}
}

func (t *Tasks) WaitFinish() {
	close(t.taskpool)
	t.wg.Wait()
}

func (t *Tasks) Run() {
	for i := uint(0); i < t.tasknum; i++ {
		t.wg.Add(1)
		go func(chanIdx uint) {
			defer t.wg.Done()
			for {
				select {
				case task, ok := <-t.taskpool:
					if !ok {
						return
					}

					t.handler(task, chanIdx)
				}
			}
		}(i)
	}
}

func DoMultiTasks(routine int, handler TaskHandle, items ...interface{}) {
	tasks := NewTasks(uint(routine), handler)
	tasks.Run()
	for _, item := range items {
		tasks.AddTask(item)
	}
	tasks.WaitFinish()
}
