package orchestration

import (
	"context"
	"fmt"
)

// Wait 阻塞到 ids 指名的每个子代理都进入终结态（或 ctx 被取消），返回且仅返回这些
// 子代理的合并结果，顺序与给定 id 的顺序一致。
//
// 与覆盖整个调度器的 Join 不同，Wait 把等待与合并都限定在某个调用方自己的批次上，
// 因此不同会话的并发派发绝不会互相等待、也不会合并到彼此的子代理。派发工具正是用
// 它收集自己即将交回的结果。
//
// 未知 id 返回 ErrAgentNotFound（以 %w 包装）；ctx 先行取消时返回包装后的 ctx 错误
// 与部分合并结果。
func (s *Scheduler) Wait(ctx context.Context, ids []string) (*MergedResult, error) {
	if ctx == nil {
		ctx = context.Background()
	}
	subs, err := s.pick(ids)
	if err != nil {
		return nil, err
	}
	for _, sub := range subs {
		select {
		case <-sub.waitDone():
		case <-ctx.Done():
			return mergeSubs(subs), fmt.Errorf("orchestration: wait: %w", ctx.Err())
		}
	}
	return mergeSubs(subs), nil
}

// pick 按给定顺序把 ids 解析为已登记的子代理。遇到第一个未知 id 时返回包装后的
// ErrAgentNotFound。
func (s *Scheduler) pick(ids []string) ([]*SubAgent, error) {
	subs := make([]*SubAgent, 0, len(ids))
	s.mu.Lock()
	defer s.mu.Unlock()
	for _, id := range ids {
		sub, ok := s.agents[id]
		if !ok {
			return nil, fmt.Errorf("orchestration: wait %q: %w", id, ErrAgentNotFound)
		}
		subs = append(subs, sub)
	}
	return subs, nil
}

// merge 把所有已登记的子代理聚合成一个 MergedResult。它由 Join 在全部子代理终结后
// 调用。
func (s *Scheduler) merge() *MergedResult {
	s.mu.Lock()
	subs := make([]*SubAgent, 0, len(s.agents))
	for _, sub := range s.agents {
		subs = append(subs, sub)
	}
	s.mu.Unlock()
	return mergeSubs(subs)
}

// mergeSubs 把给定子代理聚合成一个 MergedResult：统计终结态数量，并按切片顺序追加
// 各自的快照。仍在运行的子代理贡献其当前快照，且不计入 Succeeded/Failed/Aborted
// 中任何一项。
func mergeSubs(subs []*SubAgent) *MergedResult {
	res := &MergedResult{Total: len(subs), Agents: make([]SubAgentResult, 0, len(subs))}
	for _, sub := range subs {
		snap := sub.Snapshot()
		res.Agents = append(res.Agents, snap)
		switch snap.Status {
		case StatusDone:
			res.Succeeded++
		case StatusFailed:
			res.Failed++
		case StatusAborted:
			res.Aborted++
		}
	}
	res.Summary = summarizeMerged(res)
	return res
}