package worker

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

// Job 任务接口，所有后台任务需实现此接口
type Job interface {
	Execute(ctx context.Context) error
	JobID() string
}

// TaskResult 任务执行结果
type TaskResult struct {
	JobID  string
	Output any
	Error  error
}

// Pool 任务并发池，管理固定数量工作协程
type Pool struct {
	jobs     chan Job
	results  chan TaskResult
	wg       sync.WaitGroup
	stopOnce sync.Once
	ctx      context.Context
	cancel   context.CancelFunc
	running  atomic.Int64
	total    atomic.Int64
	failed   atomic.Int64
}

// Options 池配置
type Options struct {
	Workers int // 工作协程数，默认 4
	BufSize int // 队列缓冲，默认 1000
}

// NewPool 创建任务池
func NewPool(opts ...Options) *Pool {
	o := Options{Workers: 4, BufSize: 1000}
	if len(opts) > 0 {
		o = opts[0]
	}
	if o.Workers <= 0 {
		o.Workers = 4
	}
	if o.BufSize <= 0 {
		o.BufSize = 1000
	}
	ctx, cancel := context.WithCancel(context.Background())
	p := &Pool{
		jobs:    make(chan Job, o.BufSize),
		results: make(chan TaskResult, o.BufSize),
		ctx:     ctx,
		cancel:  cancel,
	}
	p.startWorkers()
	return p
}

func (p *Pool) startWorkers() {
	for i := 0; i < 4; i++ {
		p.wg.Add(1)
		go p.worker()
	}
}

func (p *Pool) worker() {
	defer p.wg.Done()
	for {
		select {
		case <-p.ctx.Done():
			return
		case job, ok := <-p.jobs:
			if !ok {
				return
			}
			p.running.Add(1)
			p.total.Add(1)
			err := job.Execute(p.ctx)
			result := TaskResult{JobID: job.JobID(), Error: err}
			if err != nil {
				p.failed.Add(1)
			}
			p.running.Add(-1)
			select {
			case p.results <- result:
			case <-p.ctx.Done():
				return
			}
		}
	}
}

// Submit 提交任务到队列（非阻塞，队列满时返回 false）
func (p *Pool) Submit(job Job) bool {
	select {
	case p.jobs <- job:
		return true
	default:
		return false
	}
}

// SubmitAndWait 提交任务并等待结果
func (p *Pool) SubmitAndWait(ctx context.Context, job Job) (*TaskResult, error) {
	if !p.Submit(job) {
		return nil, ErrQueueFull
	}
	select {
	case r := <-p.results:
		return &r, r.Error
	case <-ctx.Done():
		return nil, ctx.Err()
	}
}

// Results 返回结果通道，供外部 range 消费
func (p *Pool) Results() <-chan TaskResult {
	return p.results
}

// Stats 当前池状态
func (p *Pool) Stats() (running, total, failed int64) {
	return p.running.Load(), p.total.Load(), p.failed.Load()
}

// Stop 优雅停止池
func (p *Pool) Stop() {
	p.stopOnce.Do(func() {
		p.cancel()
		close(p.jobs)
		p.wg.Wait()
		close(p.results)
	})
}

var ErrQueueFull = SyncError("任务队列已满，请降低提交频率或增大 BufSize")

// SyncError 用于类型断言
type SyncError string

func (e SyncError) Error() string { return string(e) }