package concurrency

import (
	"errors"
	"fmt"
	"sync"
	"time"

	"github.com/pingcap/dbaas/billing2/utils/log"
	"github.com/pingcap/dbaas/billing2/utils/timeutils"
)

type Pool struct {
	capacity uint
	queue    chan func()
	wg       sync.WaitGroup
	now      *time.Time
}

func panicWrapper(f func()) func() {
	return func() {
		defer func() {
			if err := recover(); err != nil {
				log.Error(fmt.Sprintf("pool run time panic: %v", err))
			}
		}()
		f()
	}
}

func (p *Pool) SetNow(time time.Time) {
	p.now = &time
}

func (p *Pool) GetNow() time.Time {
	return *p.now
}

func NewPool(capacity uint) (*Pool, error) {
	if capacity == 0 {
		return nil, errors.New("capacity should be positive")
	}
	return &Pool{
		capacity: capacity,
		queue:    make(chan func()),
	}, nil
}

func (p *Pool) Start() {
	for i := uint(0); i != p.capacity; i++ {
		go p.workerRoutine()
	}
}

func (p *Pool) workerRoutine() {
	for task := range p.queue {
		if p.now != nil {
			timeutils.SetMockNow(p.GetNow())
		}
		task = panicWrapper(task)
		task()
		p.wg.Done()
	}
}

func (p *Pool) Submit(task func()) {
	p.wg.Add(1)
	p.queue <- task
}

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