package workflow

import (
	"agent/src/common"
	"agent/src/llm"
	"agent/src/utils/log"
	"context"
	"errors"
	"sync"
	"time"
)

const MaxPendingPerKey_Local  = 5
const MaxPendingPerKey_Remote = 100
const PendingTimeOut = 10*time.Minute

var TaskPools map[string]*TaskPool
var once sync.Once

type TaskPool struct {
	ctx         context.Context
	queue       *taskQueue
	llmGroup    *llm.LLMGroup
	sem         chan struct{} // 控制task并发数量
}

func InitTaskPool(ctx context.Context) {
	once.Do(func(){
		TaskPools = make(map[string]*TaskPool)
		for model, g := range llm.LLMGroups {
			pool := newTaskPool(ctx, g)
			TaskPools[model] = pool
			pool.startloop()
		}
	})
}

func newTaskPool(ctx context.Context, g *llm.LLMGroup) *TaskPool {
	var queueCap, workers int
	if g.IsLocal() {
		queueCap = 20*g.Cap()
		workers = 20
	} else {
		workers = 500
		queueCap = 10_000*g.Cap()
	}
	return &TaskPool{
		ctx     : ctx,
		queue   : newTaskQueue(queueCap),
		llmGroup: g,
		sem     :  make(chan struct{}, min(workers, 500)),
	}
}


func (t *TaskPool) add(task *Task) error {
	ok, evict := t.queue.Push(task)
	if !ok {
		task.SendExecFail("agent too busy")
		return errors.New("agent too busy")
	}
	if evict != nil {
		evict.SendExecFail("cancelled")
	}
	task.pool = t
	return nil
}


func (t *TaskPool) remove(client common.WsClient) {
	t.queue.Remove(client)
}


func (t *TaskPool) startloop() {
	done := t.ctx.Done()
	tick := time.NewTicker(PendingTimeOut)
	go func(){
		var task *Task
		for {
			task = t.queue.WaitForNew(t.ctx)
			select {
				case <-done: return
				case <-task.ctx.Done():
					task.pool.remove(task.client)
				case t.sem <- struct{}{}:
					go task.do()
			}
			
		}
	}()
	go func(){
		for {
			select {
				case <-done: return
				case <-tick.C:
					t.stat()
				default:
			}
		}
	}()
}

func (t *TaskPool) stat() {
	// 清除pending太久的task
	t.queue.RemoveTimeout(PendingTimeOut)
	log.Info("[TaskPool] stat", "queueLen", t.queue.Len())
}
