package worker

import (
	"context"
	"fmt"
	"log"
	"sync"
)

// Task 定义通用任务接口
type Task interface {
	Execute() (interface{}, error)
	ID() string
}

// Result 任务执行结果
type Result struct {
	TaskID  string
	Data    interface{}
	Err     error
}

// Pool 通用 Worker Pool，控制最大并发数
type Pool struct {
	maxWorkers int
	tasks      chan Task
	results    chan Result
	wg         sync.WaitGroup
	ctx        context.Context
	cancel     context.CancelFunc
}

// NewPool 创建 Worker Pool
// maxWorkers 为最大并发数，防止设备过载
func NewPool(maxWorkers int, queueSize int) *Pool {
	if maxWorkers <= 0 {
		maxWorkers = 1
	}
	if queueSize <= 0 {
		queueSize = maxWorkers * 2
	}

	ctx, cancel := context.WithCancel(context.Background())

	return &Pool{
		maxWorkers: maxWorkers,
		tasks:      make(chan Task, queueSize),
		results:    make(chan Result, queueSize),
		ctx:        ctx,
		cancel:     cancel,
	}
}

// Start 启动 Worker Pool
func (p *Pool) Start() {
	for i := 0; i < p.maxWorkers; i++ {
		p.wg.Add(1)
		go p.worker(i)
	}
}

// worker 单个工作协程
func (p *Pool) worker(id int) {
	defer p.wg.Done()

	for {
		select {
		case <-p.ctx.Done():
			return
		case task, ok := <-p.tasks:
			if !ok {
				return
			}
			data, err := task.Execute()
			result := Result{
				TaskID: task.ID(),
				Data:   data,
				Err:    err,
			}

			select {
			case p.results <- result:
			case <-p.ctx.Done():
				return
			}
		}
	}
}

// Submit 提交任务到工作池
func (p *Pool) Submit(task Task) error {
	select {
	case p.tasks <- task:
		return nil
	case <-p.ctx.Done():
		return fmt.Errorf("worker pool 已停止")
	}
}

// Results 返回结果通道
func (p *Pool) Results() <-chan Result {
	return p.results
}

// WaitDone 等待所有任务完成并关闭结果通道
func (p *Pool) WaitDone() {
	close(p.tasks)
	p.wg.Wait()
	close(p.results)
}

// Stop 停止 Worker Pool
func (p *Pool) Stop() {
	p.cancel()
}

// RunTasks 批量运行任务并收集结果
// 这是一个便捷方法，提交所有任务后等待全部完成
func (p *Pool) RunTasks(tasks []Task) []Result {
	p.Start()

	// 提交所有任务
	go func() {
		for _, task := range tasks {
			if err := p.Submit(task); err != nil {
				log.Printf("[WARN] 提交任务 %s 失败: %v", task.ID(), err)
				// 直接生成错误结果
				p.results <- Result{
					TaskID: task.ID(),
					Err:    err,
				}
			}
		}
	}()

	// 收集结果
	var results []Result
	for result := range p.results {
		results = append(results, result)
	}

	return results
}

// GenericTask 通用任务包装器
type GenericTask struct {
	TaskID string
	Fn     func() (interface{}, error)
}

func (t *GenericTask) ID() string { return t.TaskID }
func (t *GenericTask) Execute() (interface{}, error) {
	return t.Fn()
}
