package utils

import (
	"context"

	"github.com/gammazero/workerpool"
	"github.com/tidbcloud/dataflow-service-ng/pkg/log"
)

const (
	defaultConcurrency  = 4
	defaultMaxQueueSize = 500
)

type WorkerPoolWrapper interface {
	Submit(ctx context.Context, task func())
	SubmitWait(ctx context.Context, task func())
	StopWait()
}

type workerPoolWrapperImpl struct {
	concurrency  int
	maxQueueSize int

	wp *workerpool.WorkerPool
}

func NewWorkerPoolWrapper(ops ...WorkerPoolWrapperOptionOp) WorkerPoolWrapper {
	p := &workerPoolWrapperImpl{
		concurrency:  defaultConcurrency,
		maxQueueSize: defaultMaxQueueSize,
	}
	for _, op := range ops {
		op(p)
	}
	p.wp = workerpool.New(p.concurrency)
	return p
}

type WorkerPoolWrapperOptionOp func(*workerPoolWrapperImpl)

func WorkerPoolConcurrencyOp(concurrency int) WorkerPoolWrapperOptionOp {
	return func(p *workerPoolWrapperImpl) {
		if concurrency > 0 {
			p.concurrency = concurrency
		}
	}
}

func WorkerPoolMaxQueueSizeOp(maxQueueSize int) WorkerPoolWrapperOptionOp {
	return func(p *workerPoolWrapperImpl) {
		if maxQueueSize > 0 {
			p.maxQueueSize = maxQueueSize
		}
	}
}

func (impl *workerPoolWrapperImpl) StopWait() {
	impl.wp.StopWait()
}

func (impl *workerPoolWrapperImpl) Submit(ctx context.Context, task func()) {
	if impl.wp.WaitingQueueSize() > impl.maxQueueSize {
		impl.wp.SubmitWait(task)
		log.GetLogger(ctx).Info("submit done")
		return
	}
	impl.wp.Submit(task)
	log.GetLogger(ctx).Info("submit done")
}

func (impl *workerPoolWrapperImpl) SubmitWait(ctx context.Context, task func()) {
	doneChan := make(chan struct{})
	impl.Submit(ctx, func() {
		task()
		close(doneChan)
	})
	log.GetLogger(ctx).Info("submit done")
	<-doneChan
	log.GetLogger(ctx).Info("submit wait done")
}