package orchestration

import (
	"context"
	"fmt"
)

// Wait blocks until every sub-agent named by ids reaches a terminal state (or ctx
// is canceled) and returns the merged outcome of exactly those sub-agents, in the
// order the ids were given.
//
// Unlike Join — which covers the whole scheduler — Wait scopes both the wait and
// the merge to one caller's batch, so concurrent spawns from different sessions
// never wait on, or merge, each other's sub-agents. It is what the spawn tool uses
// to collect the results it is about to hand back.
//
// It returns ErrAgentNotFound (wrapped with %w) for an unknown id, and the wrapped
// ctx error together with the partial merged outcome when ctx is canceled first.
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 resolves ids to their registered sub-agents in the given order. It returns
// ErrAgentNotFound wrapped for the first unknown id.
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 aggregates every registered sub-agent into a MergedResult. It is called by
// Join once all sub-agents are terminal.
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 aggregates the given sub-agents into a MergedResult, counting the
// terminal statuses and appending each snapshot in slice order. A sub-agent that
// is still running contributes its current snapshot and is counted in neither
// 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
}