package tcp

import (
	"context"
	"sync"

	"spark-server/internal/pkg/errs"
)

// TaskPool runs inbound business (MySQL included). EventLoop must only Submit.
type TaskPool interface {
	Submit(ctx context.Context, task func()) error
	Close()
}

type pool struct {
	tasks   chan func()
	workers int
	wg      sync.WaitGroup
	once    sync.Once
}

func NewTaskPool(workers, queue int) TaskPool {
	if workers < 1 {
		workers = 1
	}
	if queue < 1 {
		queue = 1
	}
	p := &pool{tasks: make(chan func(), queue), workers: workers}
	for i := 0; i < workers; i++ {
		p.wg.Add(1)
		go p.loop()
	}
	return p
}

func (p *pool) loop() {
	defer p.wg.Done()
	for task := range p.tasks {
		if task != nil {
			task()
		}
	}
}

func (p *pool) Submit(ctx context.Context, task func()) error {
	if task == nil {
		return errs.ErrInvalidArgument
	}
	select {
	case <-ctx.Done():
		return ctx.Err()
	case p.tasks <- task:
		return nil
	default:
		return errs.ErrBusy
	}
}

func (p *pool) Close() {
	p.once.Do(func() {
		close(p.tasks)
		p.wg.Wait()
	})
}