package mcp

import (
	"context"
	"sync"
)

// TaskTracker 记录一次 run 内 tools/call 结果里出现过的长任务（按服务端
// 客户端分组），供 run 结束时对这些任务发 tasks/cancel（2026-07-28 最小
// 产品映射：取消只覆盖本 run 自己的任务，不影响共享 Manager 上其他 run）。
type TaskTracker struct {
	mu   sync.Mutex
	seen map[Client]map[string]struct{}
}

type trackerKey struct{}

// WithTaskTracker 把任务跟踪器放进上下文（bridgeTool 记录 TaskID）。
func WithTaskTracker(ctx context.Context) (context.Context, *TaskTracker) {
	tr := &TaskTracker{seen: map[Client]map[string]struct{}{}}
	return context.WithValue(ctx, trackerKey{}, tr), tr
}

func trackerFrom(ctx context.Context) *TaskTracker {
	if ctx == nil {
		return nil
	}
	tr, _ := ctx.Value(trackerKey{}).(*TaskTracker)
	return tr
}

func (t *TaskTracker) note(c Client, taskID string) {
	if t == nil || c == nil || taskID == "" {
		return
	}
	t.mu.Lock()
	set := t.seen[c]
	if set == nil {
		set = map[string]struct{}{}
		t.seen[c] = set
	}
	set[taskID] = struct{}{}
	t.mu.Unlock()
}

// CancelAll 对记录的任务逐一向其服务端发 tasks/cancel（尽力而为，
// 首个错误后继续）。返回实际发出取消的任务数。
func (t *TaskTracker) CancelAll(ctx context.Context) int {
	if t == nil {
		return 0
	}
	t.mu.Lock()
	seen := t.seen
	t.seen = map[Client]map[string]struct{}{}
	t.mu.Unlock()
	n := 0
	for c, ids := range seen {
		for id := range ids {
			if err := c.TasksCancel(ctx, id); err == nil {
				n++
			}
		}
	}
	return n
}