package utils

import (
	"context"
	"sync"
	"sync/atomic"

	"github.com/gogf/gf/v2/os/grpool"
)

var (
	defaultPool atomic.Pointer[grpool.Pool]
	once        sync.Once
	wg          sync.WaitGroup
)

const DefaultWorkerNum = 100

// Init 初始化全局协程池，首次调用生效，后续调用忽略。
func Init(workerNum int) {
	once.Do(func() {
		if workerNum <= 0 {
			workerNum = DefaultWorkerNum
		}
		defaultPool.Store(grpool.New(workerNum))
	})
}

// Submit 提交异步任务，上下文透传至 grpool。
// Submit 也可在 Init 前调用（自动 Init），但 Shutdown 后返回 ErrPoolClosed。
func Submit(ctx context.Context, task func(ctx context.Context)) error {
	p := defaultPool.Load()
	if p == nil {
		Init(DefaultWorkerNum)
		p = defaultPool.Load()
	}

	wg.Add(1)
	err := p.AddWithRecover(ctx, func(ctx context.Context) {
		defer wg.Done()
		task(ctx)
	}, nil)
	if err != nil {
		wg.Done()
		return err
	}
	return nil
}

// Shutdown 优雅关闭：停止新任务，等待全部已完成/排队任务完成。
func Shutdown() {
	p := defaultPool.Swap(nil)
	if p == nil {
		return
	}
	wg.Wait()
	p.Close()
}