/*
 * Copyright 2025 CloudWeGo Authors
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package adk

import (
	"context"
	"errors"
	"fmt"
	"log"
	"runtime/debug"
	"strings"

	"github.com/cloudwego/eino/callbacks"
	"github.com/cloudwego/eino/compose"
	icb "github.com/cloudwego/eino/internal/callbacks"
	"github.com/cloudwego/eino/internal/safe"
	"github.com/cloudwego/eino/schema"
)

type HistoryEntry struct {
	IsUserInput bool
	AgentName   string
	Message     Message
}

type HistoryRewriter func(ctx context.Context, entries []*HistoryEntry) ([]Message, error)

type flowAgent struct {
	Agent

	subAgents   []*flowAgent
	parentAgent *flowAgent

	disallowTransferToParent bool
	historyRewriter          HistoryRewriter

	checkPointStore compose.CheckPointStore
}

func (a *flowAgent) deepCopy() *flowAgent {
	ret := &flowAgent{
		Agent:                    a.Agent,
		subAgents:                make([]*flowAgent, 0, len(a.subAgents)),
		parentAgent:              a.parentAgent,
		disallowTransferToParent: a.disallowTransferToParent,
		historyRewriter:          a.historyRewriter,
		checkPointStore:          a.checkPointStore,
	}

	for _, sa := range a.subAgents {
		ret.subAgents = append(ret.subAgents, sa.deepCopy())
	}
	return ret
}

// SetSubAgents sets sub-agents for the given agent and returns the updated agent.
//
// NOT RECOMMENDED: Agent transfer with full context sharing between agents has not proven
// to be more effective empirically. Consider using ChatModelAgent with AgentTool
// or DeepAgent instead for most multi-agent scenarios.
func SetSubAgents(ctx context.Context, agent Agent, subAgents []Agent) (ResumableAgent, error) {
	return setSubAgents(ctx, agent, subAgents)
}

type AgentOption func(options *flowAgent)

// WithDisallowTransferToParent prevents a sub-agent from transferring to its parent.
//
// NOT RECOMMENDED: Agent transfer with full context sharing between agents has not proven
// to be more effective empirically. Consider using ChatModelAgent with AgentTool
// or DeepAgent instead for most multi-agent scenarios.
func WithDisallowTransferToParent() AgentOption {
	return func(fa *flowAgent) {
		fa.disallowTransferToParent = true
	}
}

// WithHistoryRewriter sets a rewriter to transform conversation history
// during agent transfers.
//
// NOT RECOMMENDED: Agent transfer with full context sharing between agents has not proven
// to be more effective empirically. Consider using ChatModelAgent with AgentTool
// or DeepAgent instead for most multi-agent scenarios.
func WithHistoryRewriter(h HistoryRewriter) AgentOption {
	return func(fa *flowAgent) {
		fa.historyRewriter = h
	}
}

func toFlowAgent(ctx context.Context, agent Agent, opts ...AgentOption) *flowAgent {
	var fa *flowAgent
	var ok bool
	if fa, ok = agent.(*flowAgent); !ok {
		fa = &flowAgent{Agent: agent}
	} else {
		fa = fa.deepCopy()
	}
	for _, opt := range opts {
		opt(fa)
	}

	if fa.historyRewriter == nil {
		fa.historyRewriter = buildDefaultHistoryRewriter(agent.Name(ctx))
	}

	return fa
}

// AgentWithOptions wraps an agent with flow-specific options and returns it.
//
// NOT RECOMMENDED: Agent transfer with full context sharing between agents has not proven
// to be more effective empirically. Consider using ChatModelAgent with AgentTool
// or DeepAgent instead for most multi-agent scenarios.
func AgentWithOptions(ctx context.Context, agent Agent, opts ...AgentOption) Agent {
	return toFlowAgent(ctx, agent, opts...)
}

func setSubAgents(ctx context.Context, agent Agent, subAgents []Agent) (*flowAgent, error) {
	fa := toFlowAgent(ctx, agent)

	if len(fa.subAgents) > 0 {
		return nil, errors.New("agent's sub-agents has already been set")
	}

	if onAgent, ok_ := fa.Agent.(OnSubAgents); ok_ {
		err := onAgent.OnSetSubAgents(ctx, subAgents)
		if err != nil {
			return nil, err
		}
	}

	for _, s := range subAgents {
		fsa := toFlowAgent(ctx, s)

		if fsa.parentAgent != nil {
			return nil, errors.New("agent has already been set as a sub-agent of another agent")
		}

		fsa.parentAgent = fa
		if onAgent, ok__ := fsa.Agent.(OnSubAgents); ok__ {
			err := onAgent.OnSetAsSubAgent(ctx, agent)
			if err != nil {
				return nil, err
			}

			if fsa.disallowTransferToParent {
				err = onAgent.OnDisallowTransferToParent(ctx)
				if err != nil {
					return nil, err
				}
			}
		}

		fa.subAgents = append(fa.subAgents, fsa)
	}

	return fa, nil
}

func (a *flowAgent) getAgent(ctx context.Context, name string) *flowAgent {
	for _, subAgent := range a.subAgents {
		if subAgent.Name(ctx) == name {
			return subAgent
		}
	}

	if a.parentAgent != nil && a.parentAgent.Name(ctx) == name {
		return a.parentAgent
	}

	return nil
}

func rewriteMessage(msg Message, agentName string) Message {
	var sb strings.Builder
	sb.WriteString("For context:")
	if msg.Role == schema.Assistant {
		if msg.Content != "" {
			sb.WriteString(fmt.Sprintf(" [%s] said: %s.", agentName, msg.Content))
		}
		if len(msg.ToolCalls) > 0 {
			for i := range msg.ToolCalls {
				f := msg.ToolCalls[i].Function
				sb.WriteString(fmt.Sprintf(" [%s] called tool: `%s` with arguments: %s.",
					agentName, f.Name, f.Arguments))
			}
		}
	} else if msg.Role == schema.Tool && msg.Content != "" {
		sb.WriteString(fmt.Sprintf(" [%s] `%s` tool returned result: %s.",
			agentName, msg.ToolName, msg.Content))
	}

	rewritten := schema.UserMessage(sb.String())
	if msg.MultiContent != nil {
		rewritten.MultiContent = append([]schema.ChatMessagePart{}, msg.MultiContent...) //nolint:staticcheck // backward compat with deprecated MultiContent field
	}
	if msg.UserInputMultiContent != nil {
		rewritten.UserInputMultiContent = append([]schema.MessageInputPart{}, msg.UserInputMultiContent...)
	}

	// Convert AssistantGenMultiContent to UserInputMultiContent, since the role changes to User.
	// Reasoning parts have no user input equivalent and are dropped.
	for _, part := range msg.AssistantGenMultiContent {
		switch part.Type {
		case schema.ChatMessagePartTypeText:
			rewritten.UserInputMultiContent = append(rewritten.UserInputMultiContent, schema.MessageInputPart{
				Type:  part.Type,
				Text:  part.Text,
				Extra: part.Extra,
			})
		case schema.ChatMessagePartTypeImageURL:
			if part.Image != nil {
				rewritten.UserInputMultiContent = append(rewritten.UserInputMultiContent, schema.MessageInputPart{
					Type:  part.Type,
					Image: &schema.MessageInputImage{MessagePartCommon: part.Image.MessagePartCommon},
					Extra: part.Extra,
				})
			}
		case schema.ChatMessagePartTypeAudioURL:
			if part.Audio != nil {
				rewritten.UserInputMultiContent = append(rewritten.UserInputMultiContent, schema.MessageInputPart{
					Type:  part.Type,
					Audio: &schema.MessageInputAudio{MessagePartCommon: part.Audio.MessagePartCommon},
					Extra: part.Extra,
				})
			}
		case schema.ChatMessagePartTypeVideoURL:
			if part.Video != nil {
				rewritten.UserInputMultiContent = append(rewritten.UserInputMultiContent, schema.MessageInputPart{
					Type:  part.Type,
					Video: &schema.MessageInputVideo{MessagePartCommon: part.Video.MessagePartCommon},
					Extra: part.Extra,
				})
			}
		}
	}

	return rewritten
}

func genMsg(entry *HistoryEntry, agentName string) (Message, error) {
	msg := entry.Message
	if entry.AgentName != agentName {
		msg = rewriteMessage(msg, entry.AgentName)
	}

	return msg, nil
}

func deepCopyAgentInput(ai *AgentInput) *AgentInput {
	copied := &AgentInput{
		Messages:        make([]Message, len(ai.Messages)),
		EnableStreaming: ai.EnableStreaming,
	}

	copy(copied.Messages, ai.Messages)

	return copied
}

func (a *flowAgent) genAgentInput(ctx context.Context, runCtx *runContext, skipTransferMessages bool) (*AgentInput, error) {
	input := deepCopyAgentInput(runCtx.RootInput)

	events := runCtx.Session.getEvents()
	historyEntries := make([]*HistoryEntry, 0)

	for _, m := range input.Messages {
		historyEntries = append(historyEntries, &HistoryEntry{
			IsUserInput: true,
			Message:     m,
		})
	}

	for _, event := range events {
		if skipTransferMessages && event.Action != nil && event.Action.TransferToAgent != nil {
			// If skipTransferMessages is true and the event contain transfer action, the message in this event won't be appended to history entries.
			if event.Output != nil &&
				event.Output.MessageOutput != nil &&
				event.Output.MessageOutput.Role == schema.Tool &&
				len(historyEntries) > 0 {
				// If the skipped message's role is Tool, remove the previous history entry as it's also a transfer message(from ChatModelAgent and GenTransferMessages).
				historyEntries = historyEntries[:len(historyEntries)-1]
			}
			continue
		}

		msg, err := getMessageFromWrappedEvent(event)
		if err != nil {
			var retryErr *WillRetryError
			if errors.As(err, &retryErr) {
				log.Printf("failed to get message from event, but will retry: %v", err)
				continue
			}
			return nil, err
		}

		if msg == nil {
			continue
		}

		historyEntries = append(historyEntries, &HistoryEntry{
			AgentName: event.AgentName,
			Message:   msg,
		})
	}

	messages, err := a.historyRewriter(ctx, historyEntries)
	if err != nil {
		return nil, err
	}
	input.Messages = messages

	return input, nil
}

func buildDefaultHistoryRewriter(agentName string) HistoryRewriter {
	return func(ctx context.Context, entries []*HistoryEntry) ([]Message, error) {
		messages := make([]Message, 0, len(entries))
		var err error
		for _, entry := range entries {
			msg := entry.Message
			if !entry.IsUserInput {
				msg, err = genMsg(entry, agentName)
				if err != nil {
					return nil, fmt.Errorf("gen agent input failed: %w", err)
				}
			}

			if msg != nil {
				messages = append(messages, msg)
			}
		}

		return messages, nil
	}
}

func (a *flowAgent) Run(ctx context.Context, input *AgentInput, opts ...AgentRunOption) *AsyncIterator[*AgentEvent] {
	agentName := a.Name(ctx)

	var runCtx *runContext
	ctx, runCtx = initRunCtx(ctx, agentName, input)
	ctx = AppendAddressSegment(ctx, AddressSegmentAgent, agentName)

	o := getCommonOptions(nil, opts...)
	cancelCtx := o.cancelCtx

	processedInput, err := a.genAgentInput(ctx, runCtx, o.skipTransferMessages)
	if err != nil {
		if cancelCtx != nil {
			cancelCtx.markDone()
		}
		cbInput := &AgentCallbackInput{Input: input}
		ctx = callbacks.OnStart(ctx, cbInput)
		return wrapIterWithOnEnd(ctx, genErrorIter(err))
	}

	ctxForSubAgents := ctx

	agentType := getAgentType(a.Agent)
	ctx = initAgentCallbacks(ctx, agentName, agentType, filterOptions(agentName, opts)...)
	cbInput := &AgentCallbackInput{Input: processedInput}
	ctx = callbacks.OnStart(ctx, cbInput)

	input = processedInput

	if wf, ok := a.Agent.(*workflowAgent); ok {
		ctx = withCancelContext(ctx, cancelCtx)
		filteredOpts := filterCancelOption(filterCallbackHandlersForNestedAgents(agentName, opts))
		iter := wf.Run(ctx, input, filteredOpts...)
		iter = wrapIterWithCancelCtx(iter, cancelCtx)
		return wrapIterWithOnEnd(ctx, iter)
	}

	aIter := a.Agent.Run(withCancelContext(ctx, cancelCtx), input, filterOptions(agentName, opts)...)

	iterator, generator := NewAsyncIteratorPair[*AgentEvent]()

	go a.run(withCancelContext(ctx, cancelCtx), withCancelContext(ctxForSubAgents, cancelCtx), runCtx, aIter, generator, filterCancelOption(opts)...)

	return wrapIterWithCancelCtx(iterator, cancelCtx)
}

func (a *flowAgent) Resume(ctx context.Context, info *ResumeInfo, opts ...AgentRunOption) *AsyncIterator[*AgentEvent] {
	agentName := a.Name(ctx)

	ctx, info = buildResumeInfo(ctx, agentName, info)

	ctxForSubAgents := ctx

	o := getCommonOptions(nil, opts...)
	cancelCtx := o.cancelCtx

	agentType := getAgentType(a.Agent)
	ctx = initAgentCallbacks(ctx, agentName, agentType, filterOptions(agentName, opts)...)
	cbInput := &AgentCallbackInput{ResumeInfo: info}
	ctx = callbacks.OnStart(ctx, cbInput)

	if info.WasInterrupted {
		if ra, ok := a.Agent.(ResumableAgent); ok {
			if _, ok := ra.(*workflowAgent); ok {
				ctx = withCancelContext(ctx, cancelCtx)
				filteredOpts := filterCancelOption(filterCallbackHandlersForNestedAgents(agentName, opts))
				aIter := ra.Resume(ctx, info, filteredOpts...)
				aIter = wrapIterWithCancelCtx(aIter, cancelCtx)
				return wrapIterWithOnEnd(ctx, aIter)
			}

			aIter := ra.Resume(withCancelContext(ctx, cancelCtx), info, opts...)

			iterator, generator := NewAsyncIteratorPair[*AgentEvent]()
			go a.run(withCancelContext(ctx, cancelCtx), withCancelContext(ctxForSubAgents, cancelCtx), getRunCtx(ctxForSubAgents), aIter, generator, filterCancelOption(opts)...)
			return wrapIterWithCancelCtx(iterator, cancelCtx)
		}

		if cancelCtx != nil {
			cancelCtx.markDone()
		}
		return wrapIterWithOnEnd(ctx, genErrorIter(fmt.Errorf("failed to resume agent: agent '%s' is an interrupt point "+
			"but is not a ResumableAgent", agentName)))
	}

	nextAgentName, err := getNextResumeAgent(ctx, info)
	if err != nil {
		if cancelCtx != nil {
			cancelCtx.markDone()
		}
		return wrapIterWithOnEnd(ctx, genErrorIter(err))
	}

	subAgent := a.getAgent(ctxForSubAgents, nextAgentName)
	if subAgent == nil {
		if len(a.subAgents) == 0 {
			if ra, ok := a.Agent.(ResumableAgent); ok {
				ctx = withCancelContext(ctx, cancelCtx)
				innerIter := ra.Resume(ctx, info, filterCancelOption(opts)...)
				return wrapIterWithCancelCtx(wrapIterWithOnEnd(ctx, innerIter), cancelCtx)
			}
			return wrapIterWithOnEnd(ctx, genErrorIter(fmt.Errorf(
				"failed to resume agent: agent '%s' (type %T) has no sub-agents and does not implement ResumableAgent interface. "+
					"To support resume, your custom agent wrapper must implement the ResumableAgent interface", agentName, a.Agent)))
		}
		if cancelCtx != nil {
			cancelCtx.markDone()
		}
		return wrapIterWithOnEnd(ctx, genErrorIter(fmt.Errorf("failed to resume agent: sub-agent '%s' not found in agent '%s'", nextAgentName, agentName)))
	}

	ctxForSubAgents = withCancelContext(ctxForSubAgents, cancelCtx)
	innerIter := subAgent.Resume(ctxForSubAgents, info, filterCancelOption(opts)...)
	return wrapIterWithCancelCtx(wrapIterWithOnEnd(ctx, innerIter), cancelCtx)
}

// DeterministicTransferConfig is the configuration for AgentWithDeterministicTransferTo.
//
// NOT RECOMMENDED: Agent transfer with full context sharing between agents has not proven
// to be more effective empirically. Consider using ChatModelAgent with AgentTool
// or DeepAgent instead for most multi-agent scenarios.
type DeterministicTransferConfig struct {
	Agent        Agent
	ToAgentNames []string
}

func (a *flowAgent) run(
	ctx context.Context,
	ctxForSubAgents context.Context,
	runCtx *runContext,
	aIter *AsyncIterator[*AgentEvent],
	generator *AsyncGenerator[*AgentEvent],
	opts ...AgentRunOption) {

	cbIter, cbGen := NewAsyncIteratorPair[*AgentEvent]()

	cbOutput := &AgentCallbackOutput{Events: cbIter}
	icb.On(ctx, cbOutput, icb.BuildOnEndHandleWithCopy(copyAgentCallbackOutput), callbacks.TimingOnEnd, false)

	defer func() {
		panicErr := recover()
		if panicErr != nil {
			e := safe.NewPanicErr(panicErr, debug.Stack())
			generator.Send(&AgentEvent{Err: e})
		}

		cbGen.Close()
		generator.Close()
	}()

	var lastAction *AgentAction
	for {
		event, ok := aIter.Next()
		if !ok {
			break
		}

		// RunPath ownership: the eino framework sets RunPath exactly once.
		// If event.RunPath is already set (e.g., by agentTool), we don't modify it.
		// If event.RunPath is nil/empty, we set it to the current runCtx.RunPath.
		// This ensures RunPath is set exactly once and not duplicated.
		if len(event.RunPath) == 0 {
			event.AgentName = a.Name(ctx)
			event.RunPath = runCtx.RunPath
		}
		// Recording policy: exact RunPath match (non-interrupt) indicates events belonging to this agent execution.
		// This prevents parent recording of child/tool-internal emissions.
		if (event.Action == nil || event.Action.Interrupted == nil) && exactRunPathMatch(runCtx.RunPath, event.RunPath) {
			// copy the event so that the copied event's stream is exclusive for any potential consumer
			// copy before adding to session because once added to session it's stream could be consumed by genAgentInput at any time
			// interrupt action are not added to session, because ALL information contained in it
			// is either presented to end-user, or made available to agents through other means
			copied := copyTypedAgentEvent(event)
			setAutomaticClose(copied)
			setAutomaticClose(event)
			runCtx.Session.addEvent(copied)
		}
		// Action gating uses exact run-path match as well:
		// only actions originating from this agent execution (not child/tool runs)
		// should influence parent control flow (exit/transfer/interrupt).
		if exactRunPathMatch(runCtx.RunPath, event.RunPath) {
			lastAction = event.Action
		}
		copied := copyTypedAgentEvent(event)
		setAutomaticClose(copied)
		setAutomaticClose(event)
		cbGen.Send(copied)
		generator.Send(event)
	}

	var destName string
	if lastAction != nil {
		if lastAction.Interrupted != nil {
			return
		}
		if lastAction.Exit {
			return
		}

		if lastAction.TransferToAgent != nil {
			destName = lastAction.TransferToAgent.DestAgentName
		}
	}

	// handle transferring to another agent
	if destName != "" {
		agentToRun := a.getAgent(ctxForSubAgents, destName)
		if agentToRun == nil {
			e := fmt.Errorf("transfer failed: agent '%s' not found when transferring from '%s'",
				destName, a.Name(ctxForSubAgents))
			generator.Send(&AgentEvent{Err: e})
			return
		}

		subAIter := agentToRun.Run(ctxForSubAgents, nil /*subagents get input from runCtx*/, opts...)
		for {
			subEvent, ok_ := subAIter.Next()
			if !ok_ {
				break
			}

			setAutomaticClose(subEvent)
			generator.Send(subEvent)
		}
	}
}

func exactRunPathMatch(aPath, bPath []RunStep) bool {
	if len(aPath) != len(bPath) {
		return false
	}
	for i := range aPath {
		if !aPath[i].Equals(bPath[i]) {
			return false
		}
	}
	return true
}

func wrapIterWithOnEnd(ctx context.Context, iter *AsyncIterator[*AgentEvent]) *AsyncIterator[*AgentEvent] {
	cbIter, cbGen := NewAsyncIteratorPair[*AgentEvent]()
	cbOutput := &AgentCallbackOutput{Events: cbIter}
	icb.On(ctx, cbOutput, icb.BuildOnEndHandleWithCopy(copyAgentCallbackOutput), callbacks.TimingOnEnd, false)

	outIter, outGen := NewAsyncIteratorPair[*AgentEvent]()
	go func() {
		defer func() {
			cbGen.Close()
			outGen.Close()
		}()
		for {
			event, ok := iter.Next()
			if !ok {
				break
			}
			copied := copyTypedAgentEvent(event)
			cbGen.Send(copied)
			outGen.Send(event)
		}
	}()
	return outIter
}

// ---------------------------------------------------------------------------
// Typed wrapper for the agentic path (TypedAgent[*schema.AgenticMessage]).
//
// typedFlowAgent is a minimal wrapper used exclusively by TypedRunner and
// AgentTool to execute a TypedAgent[*schema.AgenticMessage]. It handles
// callbacks, event recording, and run-path tracking. Transfer, sub-agent
// orchestration, and history rewriting are handled solely by the concrete
// flowAgent (the *schema.Message path).
// ---------------------------------------------------------------------------

type typedFlowAgent[M MessageType] struct {
	TypedAgent[M]

	checkPointStore compose.CheckPointStore
}

func toTypedFlowAgent[M MessageType](agent TypedAgent[M]) *typedFlowAgent[M] {
	if fa, ok := agent.(*typedFlowAgent[M]); ok {
		return fa
	}
	return &typedFlowAgent[M]{TypedAgent: agent}
}

func getTypedAgentType[M MessageType](agent TypedAgent[M]) string {
	if msgAgent, ok := any(agent).(Agent); ok {
		return getAgentType(msgAgent)
	}
	if typer, ok := any(agent).(interface{ GetType() string }); ok {
		return typer.GetType()
	}
	return ""
}

func (a *typedFlowAgent[M]) Run(ctx context.Context, input *TypedAgentInput[M], opts ...AgentRunOption) *AsyncIterator[*TypedAgentEvent[M]] {
	agentName := a.Name(ctx)

	var runCtx *runContext
	ctx, runCtx = initTypedRunCtx(ctx, agentName, input)
	ctx = AppendAddressSegment(ctx, AddressSegmentAgent, agentName)

	o := getCommonOptions(nil, opts...)
	cancelCtx := o.cancelCtx

	ctxForSubAgents := ctx

	agentType := getTypedAgentType(a.TypedAgent)
	ctx = initAgenticCallbacks(ctx, agentName, agentType, filterOptions(agentName, opts)...)
	cbInput := &TypedAgentCallbackInput[*schema.AgenticMessage]{Input: any(input).(*TypedAgentInput[*schema.AgenticMessage])}
	ctx = callbacks.OnStart(ctx, cbInput)

	aIter := a.TypedAgent.Run(withCancelContext(ctx, cancelCtx), input, filterOptions(agentName, opts)...)

	iterator, generator := NewAsyncIteratorPair[*TypedAgentEvent[M]]()

	go a.run(withCancelContext(ctx, cancelCtx), withCancelContext(ctxForSubAgents, cancelCtx), runCtx, aIter, generator, filterCancelOption(opts)...)

	return wrapIterWithCancelCtx(iterator, cancelCtx)
}

func (a *typedFlowAgent[M]) Resume(ctx context.Context, info *ResumeInfo, opts ...AgentRunOption) *AsyncIterator[*TypedAgentEvent[M]] {
	agentName := a.Name(ctx)

	ctx, info = buildResumeInfo(ctx, agentName, info)

	ctxForSubAgents := ctx

	o := getCommonOptions(nil, opts...)
	cancelCtx := o.cancelCtx

	agentType := getTypedAgentType(a.TypedAgent)
	ctx = initAgenticCallbacks(ctx, agentName, agentType, filterOptions(agentName, opts)...)
	cbInput := &TypedAgentCallbackInput[*schema.AgenticMessage]{ResumeInfo: info}
	ctx = callbacks.OnStart(ctx, cbInput)

	if info.WasInterrupted {
		if ra, ok := a.TypedAgent.(TypedResumableAgent[M]); ok {
			aIter := ra.Resume(withCancelContext(ctx, cancelCtx), info, opts...)

			iterator, generator := NewAsyncIteratorPair[*TypedAgentEvent[M]]()
			go a.run(withCancelContext(ctx, cancelCtx), withCancelContext(ctxForSubAgents, cancelCtx), getRunCtx(ctxForSubAgents), aIter, generator, filterCancelOption(opts)...)
			return wrapIterWithCancelCtx(iterator, cancelCtx)
		}

		if cancelCtx != nil {
			cancelCtx.markDone()
		}
		return typedErrorIterWithOnEnd[M](ctx, fmt.Errorf("failed to resume agent: agent '%s' is an interrupt point "+
			"but is not a ResumableAgent", agentName))
	}

	_, err := getNextResumeAgent(ctx, info)
	if err != nil {
		if cancelCtx != nil {
			cancelCtx.markDone()
		}
		return typedErrorIterWithOnEnd[M](ctx, err)
	}

	if ra, ok := a.TypedAgent.(TypedResumableAgent[M]); ok {
		ctx = withCancelContext(ctx, cancelCtx)
		innerIter := ra.Resume(ctx, info, filterCancelOption(opts)...)
		return wrapIterWithCancelCtx(typedWrapIterWithOnEnd(ctx, innerIter), cancelCtx)
	}
	return typedErrorIterWithOnEnd[M](ctx, fmt.Errorf(
		"failed to resume agent: agent '%s' (type %T) does not implement ResumableAgent interface. "+
			"To support resume, your custom agent wrapper must implement the ResumableAgent interface", agentName, a.TypedAgent))
}

func (a *typedFlowAgent[M]) run(
	ctx context.Context,
	_ context.Context,
	runCtx *runContext,
	aIter *AsyncIterator[*TypedAgentEvent[M]],
	generator *AsyncGenerator[*TypedAgentEvent[M]],
	_ ...AgentRunOption) {

	agenticCbIter, agenticCbGen := NewAsyncIteratorPair[*TypedAgentEvent[*schema.AgenticMessage]]()
	cbOutput := &TypedAgentCallbackOutput[*schema.AgenticMessage]{Events: agenticCbIter}
	icb.On(ctx, cbOutput, icb.BuildOnEndHandleWithCopy(copyTypedCallbackOutput[*schema.AgenticMessage]), callbacks.TimingOnEnd, false)

	defer func() {
		panicErr := recover()
		if panicErr != nil {
			e := safe.NewPanicErr(panicErr, debug.Stack())
			generator.Send(&TypedAgentEvent[M]{Err: e})
		}

		agenticCbGen.Close()
		generator.Close()
	}()

	for {
		event, ok := aIter.Next()
		if !ok {
			break
		}

		if len(event.RunPath) == 0 {
			event.AgentName = a.Name(ctx)
			event.RunPath = runCtx.RunPath
		}
		if (event.Action == nil || event.Action.Interrupted == nil) && exactRunPathMatch(runCtx.RunPath, event.RunPath) {
			copied := copyTypedAgentEvent(event)
			typedSetAutomaticClose(copied)
			typedSetAutomaticClose(event)
			addTypedEvent(runCtx.Session, copied)
		}

		agenticCopied := copyTypedAgentEvent(event)
		typedSetAutomaticClose(agenticCopied)
		typedSetAutomaticClose(event)
		agenticCbGen.Send(any(agenticCopied).(*TypedAgentEvent[*schema.AgenticMessage]))
		generator.Send(event)
	}
}

func wrapAgenticIterWithOnEnd(ctx context.Context, iter *AsyncIterator[*TypedAgentEvent[*schema.AgenticMessage]]) *AsyncIterator[*TypedAgentEvent[*schema.AgenticMessage]] {
	cbIter, cbGen := NewAsyncIteratorPair[*TypedAgentEvent[*schema.AgenticMessage]]()
	cbOutput := &TypedAgentCallbackOutput[*schema.AgenticMessage]{Events: cbIter}
	icb.On(ctx, cbOutput, icb.BuildOnEndHandleWithCopy(copyTypedCallbackOutput[*schema.AgenticMessage]), callbacks.TimingOnEnd, false)

	outIter, outGen := NewAsyncIteratorPair[*TypedAgentEvent[*schema.AgenticMessage]]()
	go func() {
		defer func() {
			cbGen.Close()
			outGen.Close()
		}()
		for {
			event, ok := iter.Next()
			if !ok {
				break
			}
			copied := copyTypedAgentEvent(event)
			cbGen.Send(copied)
			outGen.Send(event)
		}
	}()
	return outIter
}

func genAgenticErrorIter(err error) *AsyncIterator[*TypedAgentEvent[*schema.AgenticMessage]] {
	iter, gen := NewAsyncIteratorPair[*TypedAgentEvent[*schema.AgenticMessage]]()
	gen.Send(&TypedAgentEvent[*schema.AgenticMessage]{Err: err})
	gen.Close()
	return iter
}

func typedWrapIterWithOnEnd[M MessageType](ctx context.Context, iter *AsyncIterator[*TypedAgentEvent[M]]) *AsyncIterator[*TypedAgentEvent[M]] {
	agenticIter := any(iter).(*AsyncIterator[*TypedAgentEvent[*schema.AgenticMessage]])
	return any(wrapAgenticIterWithOnEnd(ctx, agenticIter)).(*AsyncIterator[*TypedAgentEvent[M]])
}

func typedErrorIterWithOnEnd[M MessageType](ctx context.Context, err error) *AsyncIterator[*TypedAgentEvent[M]] {
	return typedWrapIterWithOnEnd(ctx, typedErrorIter[M](err))
}