package loop

import "context"

// flowBranchKey is the context key marking a Run as a branch of another
// session's flow graph. It is set by the orchestration layer (through
// WithFlowBranch) before it drives a sub-agent's Run.
type flowBranchKey struct{}

// WithFlowBranch returns ctx marking the current Run as a branch: its spans fold
// into parentSession's live flow graph and the Run neither starts nor ends a
// round of its own — the owning session's round does. It is how the orchestration
// layer attributes a sub-agent's work to the session that spawned it, so parallel
// sub-agents render as swimlanes of one graph instead of replacing the frontend's
// single flowGraph with a graph of their own (design/19 §2).
//
// An empty parentSession still marks the Run as a branch: it owns no round and
// publishes nothing. That is the state of a background sub-agent nobody attributed
// to a session, and it must not silently become a graph of its own. A nil ctx is
// returned unchanged.
func WithFlowBranch(ctx context.Context, parentSession string) context.Context {
	if ctx == nil {
		return ctx
	}
	return context.WithValue(ctx, flowBranchKey{}, parentSession)
}

// flowBranch reports the parent session recorded by WithFlowBranch and whether the
// Run is a branch at all. A non-branch Run owns its own round and publishes under
// its own session id.
func flowBranch(ctx context.Context) (string, bool) {
	if ctx == nil {
		return "", false
	}
	v := ctx.Value(flowBranchKey{})
	if v == nil {
		return "", false
	}
	s, _ := v.(string)
	return s, true
}

// flowOwner returns the session whose flow graph should receive this Run's spans:
// the branch parent when the Run is a sub-agent branch (empty when the branch is
// unattributed, which suppresses publication), otherwise the Run's own sessionID.
func flowOwner(ctx context.Context, sessionID string) string {
	if parent, isBranch := flowBranch(ctx); isBranch {
		return parent
	}
	return sessionID
}

// FlowOwnerFromContext returns the session whose flow graph owns work started
// from ctx: the branch parent when ctx runs inside a sub-agent branch, otherwise
// the session id the enclosing Run was driven with (empty when neither is set).
//
// A tool that dispatches background work — the sub-agent spawn tool — passes this
// as the parent session, so a sub-agent spawned from inside another sub-agent
// stays attributed to the same top-level graph instead of to an id that owns no
// graph.
func FlowOwnerFromContext(ctx context.Context) string {
	if parent, isBranch := flowBranch(ctx); isBranch {
		return parent
	}
	return SessionFromContext(ctx)
}