/*
 * 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 planexecute implements a plan–execute–replan style agent.
package planexecute

import (
	"context"
	"encoding/json"
	"fmt"
	"runtime/debug"
	"strings"

	"github.com/bytedance/sonic"

	"github.com/cloudwego/eino/adk"
	"github.com/cloudwego/eino/components/model"
	"github.com/cloudwego/eino/components/prompt"
	"github.com/cloudwego/eino/compose"
	"github.com/cloudwego/eino/internal/safe"
	"github.com/cloudwego/eino/schema"
)

func init() {
	schema.RegisterName[*defaultPlan]("_eino_adk_plan_execute_default_plan")
	schema.RegisterName[ExecutedStep]("_eino_adk_plan_execute_executed_step")
	schema.RegisterName[[]ExecutedStep]("_eino_adk_plan_execute_executed_steps")
}

// Plan represents an execution plan with a sequence of actionable steps.
// It supports JSON serialization and deserialization while providing access to the first step.
type Plan interface {
	// FirstStep returns the first step to be executed in the plan.
	FirstStep() string

	// Marshaler serializes the Plan into JSON.
	// The resulting JSON can be used in prompt templates.
	json.Marshaler
	// Unmarshaler deserializes JSON content into the Plan.
	// This processes output from structured chat models or tool calls into the Plan structure.
	json.Unmarshaler
}

// NewPlan is a function type that creates a new Plan instance.
type NewPlan func(ctx context.Context) Plan

// defaultPlan is the default implementation of the Plan interface.
//
// JSON Schema:
//
//	{
//	  "type": "object",
//	  "properties": {
//	    "steps": {
//	      "type": "array",
//	      "items": {
//	        "type": "string"
//	      },
//	      "description": "Ordered list of actions to be taken. Each step should be clear, actionable, and arranged in a logical sequence."
//	    }
//	  },
//	  "required": ["steps"]
//	}
type defaultPlan struct {
	// Steps contains the ordered list of actions to be taken.
	// Each step should be clear, actionable, and arranged in a logical sequence.
	Steps []string `json:"steps"`
}

// FirstStep returns the first step in the plan or an empty string if no steps exist.
func (p *defaultPlan) FirstStep() string {
	if len(p.Steps) == 0 {
		return ""
	}
	return p.Steps[0]
}

func (p *defaultPlan) MarshalJSON() ([]byte, error) {
	type planTyp defaultPlan
	return sonic.Marshal((*planTyp)(p))
}

func (p *defaultPlan) UnmarshalJSON(bytes []byte) error {
	type planTyp defaultPlan
	return sonic.Unmarshal(bytes, (*planTyp)(p))
}

// Response represents the final response to the user.
// This struct is used for JSON serialization/deserialization of the final response
// generated by the model.
type Response struct {
	// Response is the complete response to provide to the user.
	// This field is required.
	Response string `json:"response"`
}

var (
	// PlanToolInfo defines the schema for the Plan tool that can be used with ToolCallingChatModel.
	// This schema instructs the model to generate a structured plan with ordered steps.
	PlanToolInfo = schema.ToolInfo{
		Name: "plan",
		Desc: "Plan with a list of steps to execute in order. Each step should be clear, actionable, and arranged in a logical sequence. The output will be used to guide the execution process.",
		ParamsOneOf: schema.NewParamsOneOfByParams(
			map[string]*schema.ParameterInfo{
				"steps": {
					Type:     schema.Array,
					ElemInfo: &schema.ParameterInfo{Type: schema.String},
					Desc:     "different steps to follow, should be in sorted order",
					Required: true,
				},
			},
		),
	}

	// RespondToolInfo defines the schema for the response tool that can be used with ToolCallingChatModel.
	// This schema instructs the model to generate a direct response to the user.
	RespondToolInfo = schema.ToolInfo{
		Name: "respond",
		Desc: "Generate a direct response to the user. Use this tool when you have all the information needed to provide a final answer.",
		ParamsOneOf: schema.NewParamsOneOfByParams(
			map[string]*schema.ParameterInfo{
				"response": {
					Type:     schema.String,
					Desc:     "The complete response to provide to the user",
					Required: true,
				},
			},
		),
	}

	// PlannerPrompt is the prompt template for the planner.
	// It provides context and guidance to the planner on how to generate the Plan.
	PlannerPrompt = prompt.FromMessages(schema.FString,
		schema.SystemMessage(`You are an expert planning agent. Given an objective, create a comprehensive step-by-step plan to achieve the objective.

## YOUR TASK
Analyze the objective and generate a strategic plan that breaks down the goal into manageable, executable steps.

## PLANNING REQUIREMENTS
Each step in your plan must be:
- **Specific and actionable**: Clear instructions that can be executed without ambiguity
- **Self-contained**: Include all necessary context, parameters, and requirements
- **Independently executable**: Can be performed by an agent without dependencies on other steps
- **Logically sequenced**: Arranged in optimal order for efficient execution
- **Objective-focused**: Directly contribute to achieving the main goal

## PLANNING GUIDELINES
- Eliminate redundant or unnecessary steps
- Include relevant constraints, parameters, and success criteria for each step
- Ensure the final step produces a complete answer or deliverable
- Anticipate potential challenges and include mitigation strategies
- Structure steps to build upon each other logically
- Provide sufficient detail for successful execution

## QUALITY CRITERIA
- Plan completeness: Does it address all aspects of the objective?
- Step clarity: Can each step be understood and executed independently?
- Logical flow: Do steps follow a sensible progression?
- Efficiency: Is this the most direct path to the objective?
- Adaptability: Can the plan handle unexpected results or changes?`),
		schema.MessagesPlaceholder("input", false),
	)

	// ExecutorPrompt is the prompt template for the executor.
	// It provides context and guidance to the executor on how to execute the Task.
	ExecutorPrompt = prompt.FromMessages(schema.FString,
		schema.SystemMessage(`You are a diligent and meticulous executor agent. Follow the given plan and execute your tasks carefully and thoroughly.`),
		schema.UserMessage(`## OBJECTIVE
{input}
## Given the following plan:
{plan}
## COMPLETED STEPS & RESULTS
{executed_steps}
## Your task is to execute the first step, which is: 
{step}`))

	// ReplannerPrompt is the prompt template for the replanner.
	// It provides context and guidance to the replanner on how to regenerate the Plan.
	ReplannerPrompt = prompt.FromMessages(schema.FString,
		schema.SystemMessage(
			`You are going to review the progress toward an objective. Analyze the current state and determine the optimal next action.

## YOUR TASK
Based on the progress above, you MUST choose exactly ONE action:

### Option 1: COMPLETE (if objective is fully achieved)
Call '{respond_tool}' with:
- A comprehensive final answer
- Clear conclusion summarizing how the objective was met
- Key insights from the execution process

### Option 2: CONTINUE (if more work is needed)
Call '{plan_tool}' with a revised plan that:
- Contains ONLY remaining steps (exclude completed ones)
- Incorporates lessons learned from executed steps
- Addresses any gaps or issues discovered
- Maintains logical step sequence

## PLANNING REQUIREMENTS
Each step in your plan must be:
- **Specific and actionable**: Clear instructions that can be executed without ambiguity
- **Self-contained**: Include all necessary context, parameters, and requirements
- **Independently executable**: Can be performed by an agent without dependencies on other steps
- **Logically sequenced**: Arranged in optimal order for efficient execution
- **Objective-focused**: Directly contribute to achieving the main goal

## PLANNING GUIDELINES
- Eliminate redundant or unnecessary steps
- Adapt strategy based on new information
- Include relevant constraints, parameters, and success criteria for each step

## DECISION CRITERIA
- Has the original objective been completely satisfied?
- Are there any remaining requirements or sub-goals?
- Do the results suggest a need for strategy adjustment?
- What specific actions are still required?`),
		schema.UserMessage(`## OBJECTIVE
{input}

## ORIGINAL PLAN
{plan}

## COMPLETED STEPS & RESULTS
{executed_steps}`),
	)
)

const (
	// UserInputSessionKey is the session key for the user input.
	UserInputSessionKey = "UserInput"

	// PlanSessionKey is the session key for the plan.
	PlanSessionKey = "Plan"

	// ExecutedStepSessionKey is the session key for the execute result.
	ExecutedStepSessionKey = "ExecutedStep"

	// ExecutedStepsSessionKey is the session key for the execute results.
	ExecutedStepsSessionKey = "ExecutedSteps"
)

// PlannerConfig provides configuration options for creating a planner agent.
// There are two ways to configure the planner to generate structured Plan output:
//  1. Use ChatModelWithFormattedOutput: A model pre-configured to output in the Plan format
//  2. Use ToolCallingChatModel + ToolInfo: A model that uses tool calling to generate
//     the Plan structure
type PlannerConfig struct {
	// ChatModelWithFormattedOutput is a model pre-configured to output in the Plan format.
	// Create this by configuring a model to output structured data directly.
	// See example: https://github.com/cloudwego/eino-ext/blob/main/components/model/openai/examples/structured/structured.go
	ChatModelWithFormattedOutput model.BaseChatModel

	// ToolCallingChatModel is a model that supports tool calling capabilities.
	// When provided with ToolInfo, it will use tool calling to generate the Plan structure.
	ToolCallingChatModel model.ToolCallingChatModel

	// ToolInfo defines the schema for the Plan structure when using tool calling.
	// Optional. If not provided, PlanToolInfo will be used as the default.
	ToolInfo *schema.ToolInfo

	// GenInputFn is a function that generates the input messages for the planner.
	// Optional. If not provided, defaultGenPlannerInputFn will be used.
	GenInputFn GenPlannerModelInputFn

	// NewPlan creates a new Plan instance for JSON.
	// The returned Plan will be used to unmarshal the model-generated JSON output.
	// Optional. If not provided, defaultNewPlan will be used.
	NewPlan NewPlan
}

// GenPlannerModelInputFn is a function type that generates input messages for the planner.
type GenPlannerModelInputFn func(ctx context.Context, userInput []adk.Message) ([]adk.Message, error)

func defaultNewPlan(ctx context.Context) Plan {
	return &defaultPlan{}
}

func defaultGenPlannerInputFn(ctx context.Context, userInput []adk.Message) ([]adk.Message, error) {
	msgs, err := PlannerPrompt.Format(ctx, map[string]any{
		"input": userInput,
	})
	if err != nil {
		return nil, err
	}
	return msgs, nil
}

type planner struct {
	toolCall   bool
	chatModel  model.BaseChatModel
	genInputFn GenPlannerModelInputFn
	newPlan    NewPlan
}

func (p *planner) Name(_ context.Context) string {
	return "planner"
}

func (p *planner) Description(_ context.Context) string {
	return "a planner agent"
}

func argToContent(msg adk.Message) (adk.Message, error) {
	if len(msg.ToolCalls) == 0 {
		return nil, schema.ErrNoValue
	}

	return schema.AssistantMessage(msg.ToolCalls[0].Function.Arguments, nil), nil
}

func (p *planner) Run(ctx context.Context, input *adk.AgentInput,
	_ ...adk.AgentRunOption) *adk.AsyncIterator[*adk.AgentEvent] {

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

	adk.AddSessionValue(ctx, UserInputSessionKey, input.Messages)

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

			generator.Close()
		}()

		c := compose.NewChain[*adk.AgentInput, Plan]().
			AppendLambda(
				compose.InvokableLambda(func(ctx context.Context, input *adk.AgentInput) (output []adk.Message, err error) {
					return p.genInputFn(ctx, input.Messages)
				}),
			).
			AppendChatModel(p.chatModel).
			AppendLambda(
				compose.CollectableLambda(func(ctx context.Context, sr *schema.StreamReader[adk.Message]) (adk.Message, error) {
					if input.EnableStreaming {
						ss := sr.Copy(2)
						var sOutput *schema.StreamReader[*schema.Message]
						if p.toolCall {
							sOutput = schema.StreamReaderWithConvert(ss[0], argToContent)
						} else {
							sOutput = ss[0]
						}

						generator.Send(adk.EventFromMessage(nil, sOutput, schema.Assistant, ""))

						return schema.ConcatMessageStream(ss[1])
					}

					msg, err := schema.ConcatMessageStream(sr)
					if err != nil {
						return nil, err
					}

					var output adk.Message
					if p.toolCall {
						if len(msg.ToolCalls) == 0 {
							return nil, fmt.Errorf("no tool call")
						}
						output = schema.AssistantMessage(msg.ToolCalls[0].Function.Arguments, nil)
					} else {
						output = msg
					}

					generator.Send(adk.EventFromMessage(output, nil, schema.Assistant, ""))

					return msg, nil
				}),
			).
			AppendLambda(
				compose.InvokableLambda(func(ctx context.Context, msg adk.Message) (plan Plan, err error) {
					var planJSON string
					if p.toolCall {
						if len(msg.ToolCalls) == 0 {
							return nil, fmt.Errorf("no tool call")
						}
						planJSON = msg.ToolCalls[0].Function.Arguments
					} else {
						planJSON = msg.Content
					}

					plan = p.newPlan(ctx)
					err = plan.UnmarshalJSON([]byte(planJSON))
					if err != nil {
						return nil, fmt.Errorf("unmarshal plan error: %w", err)
					}

					adk.AddSessionValue(ctx, PlanSessionKey, plan)

					return plan, nil
				}),
			)

		var opts []compose.Option
		if p.toolCall {
			opts = append(opts, compose.WithChatModelOption(model.WithToolChoice(schema.ToolChoiceForced)))
		}

		r, err := c.Compile(ctx, compose.WithGraphName(p.Name(ctx)))
		if err != nil { // unexpected
			generator.Send(&adk.AgentEvent{Err: err})
			return
		}

		_, err = r.Stream(ctx, input, opts...)
		if err != nil {
			generator.Send(&adk.AgentEvent{Err: err})
			return
		}
	}()

	return iterator
}

// NewPlanner creates a new planner agent based on the provided configuration.
// The planner agent uses either ChatModelWithFormattedOutput or ToolCallingChatModel+ToolInfo
// to generate structured Plan output.
//
// If ChatModelWithFormattedOutput is provided, it will be used directly.
// If ToolCallingChatModel is provided, it will be configured with ToolInfo (or PlanToolInfo by default)
// to generate structured Plan output.
func NewPlanner(_ context.Context, cfg *PlannerConfig) (adk.Agent, error) {
	var chatModel model.BaseChatModel
	var toolCall bool
	if cfg.ChatModelWithFormattedOutput != nil {
		chatModel = cfg.ChatModelWithFormattedOutput
	} else {
		toolCall = true
		toolInfo := cfg.ToolInfo
		if toolInfo == nil {
			toolInfo = &PlanToolInfo
		}

		var err error
		chatModel, err = cfg.ToolCallingChatModel.WithTools([]*schema.ToolInfo{toolInfo})
		if err != nil {
			return nil, err
		}
	}

	inputFn := cfg.GenInputFn
	if inputFn == nil {
		inputFn = defaultGenPlannerInputFn
	}

	planParser := cfg.NewPlan
	if planParser == nil {
		planParser = defaultNewPlan
	}

	return &planner{
		toolCall:   toolCall,
		chatModel:  chatModel,
		genInputFn: inputFn,
		newPlan:    planParser,
	}, nil
}

// ExecutionContext is the input information for the executor and the planner.
type ExecutionContext struct {
	UserInput     []adk.Message
	Plan          Plan
	ExecutedSteps []ExecutedStep
}

// GenModelInputFn is a function that generates the input messages for the executor and the planner.
type GenModelInputFn func(ctx context.Context, in *ExecutionContext) ([]adk.Message, error)

// ExecutorConfig provides configuration options for creating an executor agent.
type ExecutorConfig struct {
	// Model is the chat model used by the executor.
	// If the executor uses any tools, this model must support the model.WithTools call option,
	// as that's how the executor configures the model with tool information.
	Model model.BaseChatModel

	// ToolsConfig specifies the tools available to the executor.
	ToolsConfig adk.ToolsConfig

	// MaxIterations defines the upper limit of ChatModel generation cycles.
	// The agent will terminate with an error if this limit is exceeded.
	// Optional. Defaults to 20.
	MaxIterations int

	// GenInputFn generates the input messages for the Executor.
	// Optional. If not provided, defaultGenExecutorInputFn will be used.
	GenInputFn GenModelInputFn
}

type ExecutedStep struct {
	Step   string
	Result string
}

// NewExecutor creates a new executor agent.
func NewExecutor(ctx context.Context, cfg *ExecutorConfig) (adk.Agent, error) {

	genInputFn := cfg.GenInputFn
	if genInputFn == nil {
		genInputFn = defaultGenExecutorInputFn
	}
	genInput := func(ctx context.Context, instruction string, _ *adk.AgentInput) ([]adk.Message, error) {

		plan, ok := adk.GetSessionValue(ctx, PlanSessionKey)
		if !ok {
			panic("impossible: plan not found")
		}
		plan_ := plan.(Plan)

		userInput, ok := adk.GetSessionValue(ctx, UserInputSessionKey)
		if !ok {
			panic("impossible: user input not found")
		}
		userInput_ := userInput.([]adk.Message)

		var executedSteps_ []ExecutedStep
		executedStep, ok := adk.GetSessionValue(ctx, ExecutedStepsSessionKey)
		if ok {
			executedSteps_ = executedStep.([]ExecutedStep)
		}

		in := &ExecutionContext{
			UserInput:     userInput_,
			Plan:          plan_,
			ExecutedSteps: executedSteps_,
		}

		msgs, err := genInputFn(ctx, in)
		if err != nil {
			return nil, err
		}

		return msgs, nil
	}

	agent, err := adk.NewChatModelAgent(ctx, &adk.ChatModelAgentConfig{
		Name:          "executor",
		Description:   "an executor agent",
		Model:         cfg.Model,
		ToolsConfig:   cfg.ToolsConfig,
		GenModelInput: genInput,
		MaxIterations: cfg.MaxIterations,
		OutputKey:     ExecutedStepSessionKey,
	})
	if err != nil {
		return nil, err
	}

	return agent, nil
}

func defaultGenExecutorInputFn(ctx context.Context, in *ExecutionContext) ([]adk.Message, error) {

	planContent, err := in.Plan.MarshalJSON()
	if err != nil {
		return nil, err
	}

	userMsgs, err := ExecutorPrompt.Format(ctx, map[string]any{
		"input":          formatInput(in.UserInput),
		"plan":           string(planContent),
		"executed_steps": formatExecutedSteps(in.ExecutedSteps),
		"step":           in.Plan.FirstStep(),
	})
	if err != nil {
		return nil, err
	}

	return userMsgs, nil
}

type replanner struct {
	chatModel   model.ToolCallingChatModel
	planTool    *schema.ToolInfo
	respondTool *schema.ToolInfo

	genInputFn GenModelInputFn
	newPlan    NewPlan
}

type ReplannerConfig struct {
	// ChatModel is the model that supports tool calling capabilities.
	// It will be configured with PlanTool and RespondTool to generate updated plans or responses.
	ChatModel model.ToolCallingChatModel

	// PlanTool defines the schema for the Plan tool that can be used with ToolCallingChatModel.
	// Optional. If not provided, the default PlanToolInfo will be used.
	PlanTool *schema.ToolInfo

	// RespondTool defines the schema for the response tool that can be used with ToolCallingChatModel.
	// Optional. If not provided, the default RespondToolInfo will be used.
	RespondTool *schema.ToolInfo

	// GenInputFn generates the input messages for the Replanner.
	// Optional. If not provided, buildGenReplannerInputFn will be used.
	GenInputFn GenModelInputFn

	// NewPlan creates a new Plan instance.
	// The returned Plan will be used to unmarshal the model-generated JSON output from PlanTool.
	// Optional. If not provided, defaultNewPlan will be used.
	NewPlan NewPlan
}

// formatInput formats the input messages into a string.
func formatInput(input []adk.Message) string {
	var sb strings.Builder
	for _, msg := range input {
		sb.WriteString(msg.Content)
		sb.WriteString("\n")
	}

	return sb.String()
}

func formatExecutedSteps(results []ExecutedStep) string {
	var sb strings.Builder
	for _, result := range results {
		sb.WriteString(fmt.Sprintf("Step: %s\nResult: %s\n\n", result.Step, result.Result))
	}

	return sb.String()
}

func (r *replanner) Name(_ context.Context) string {
	return "replanner"
}

func (r *replanner) Description(_ context.Context) string {
	return "a replanner agent"
}

func (r *replanner) genInput(ctx context.Context) ([]adk.Message, error) {

	executedStep, ok := adk.GetSessionValue(ctx, ExecutedStepSessionKey)
	if !ok {
		panic("impossible: execute result not found")
	}
	executedStep_ := executedStep.(string)

	plan, ok := adk.GetSessionValue(ctx, PlanSessionKey)
	if !ok {
		panic("impossible: plan not found")
	}
	plan_ := plan.(Plan)
	step := plan_.FirstStep()

	var executedSteps_ []ExecutedStep
	executedSteps, ok := adk.GetSessionValue(ctx, ExecutedStepsSessionKey)
	if ok {
		executedSteps_ = executedSteps.([]ExecutedStep)
	}

	executedSteps_ = append(executedSteps_, ExecutedStep{
		Step:   step,
		Result: executedStep_,
	})
	adk.AddSessionValue(ctx, ExecutedStepsSessionKey, executedSteps_)

	userInput, ok := adk.GetSessionValue(ctx, UserInputSessionKey)
	if !ok {
		panic("impossible: user input not found")
	}
	userInput_ := userInput.([]adk.Message)

	in := &ExecutionContext{
		UserInput:     userInput_,
		Plan:          plan_,
		ExecutedSteps: executedSteps_,
	}
	genInputFn := r.genInputFn
	if genInputFn == nil {
		genInputFn = buildGenReplannerInputFn(r.planTool.Name, r.respondTool.Name)
	}
	msgs, err := genInputFn(ctx, in)
	if err != nil {
		return nil, err
	}

	return msgs, nil
}

func (r *replanner) Run(ctx context.Context, input *adk.AgentInput, _ ...adk.AgentRunOption) *adk.AsyncIterator[*adk.AgentEvent] {
	iterator, generator := adk.NewAsyncIteratorPair[*adk.AgentEvent]()

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

			generator.Close()
		}()

		callOpt := model.WithToolChoice(schema.ToolChoiceForced)

		c := compose.NewChain[struct{}, any]().
			AppendLambda(
				compose.InvokableLambda(func(ctx context.Context, input struct{}) (output []adk.Message, err error) {
					return r.genInput(ctx)
				}),
			).
			AppendChatModel(r.chatModel).
			AppendLambda(
				compose.CollectableLambda(func(ctx context.Context, sr *schema.StreamReader[adk.Message]) (adk.Message, error) {
					if input.EnableStreaming {
						ss := sr.Copy(2)
						sOutput := schema.StreamReaderWithConvert(ss[0], argToContent)
						generator.Send(adk.EventFromMessage(nil, sOutput, schema.Assistant, ""))
						return schema.ConcatMessageStream(ss[1])
					}

					msg, err := schema.ConcatMessageStream(sr)
					if err != nil {
						return nil, err
					}
					if len(msg.ToolCalls) > 0 {
						output := schema.AssistantMessage(msg.ToolCalls[0].Function.Arguments, nil)
						generator.Send(adk.EventFromMessage(output, nil, schema.Assistant, ""))
					}
					return msg, nil
				}),
			).
			AppendLambda(
				compose.InvokableLambda(func(ctx context.Context, msg adk.Message) (msgOrPlan any, err error) {
					if len(msg.ToolCalls) == 0 {
						return nil, fmt.Errorf("no tool call")
					}

					// exit
					if msg.ToolCalls[0].Function.Name == r.respondTool.Name {
						action := adk.NewBreakLoopAction(r.Name(ctx))
						generator.Send(&adk.AgentEvent{Action: action})
						return msg, nil
					}

					// replan
					if msg.ToolCalls[0].Function.Name != r.planTool.Name {
						return nil, fmt.Errorf("unexpected tool call: %s", msg.ToolCalls[0].Function.Name)
					}

					plan := r.newPlan(ctx)
					if err = plan.UnmarshalJSON([]byte(msg.ToolCalls[0].Function.Arguments)); err != nil {
						return nil, fmt.Errorf("unmarshal plan error: %w", err)
					}

					adk.AddSessionValue(ctx, PlanSessionKey, plan)

					return plan, nil
				}),
			)

		runnable, err := c.Compile(ctx, compose.WithGraphName(r.Name(ctx)))
		if err != nil {
			generator.Send(&adk.AgentEvent{Err: err})
			return
		}

		_, err = runnable.Stream(ctx, struct{}{}, compose.WithChatModelOption(callOpt))
		if err != nil {
			generator.Send(&adk.AgentEvent{Err: err})
			return
		}
	}()

	return iterator
}

func buildGenReplannerInputFn(planToolName, respondToolName string) GenModelInputFn {
	return func(ctx context.Context, in *ExecutionContext) ([]adk.Message, error) {
		planContent, err := in.Plan.MarshalJSON()
		if err != nil {
			return nil, err
		}
		msgs, err := ReplannerPrompt.Format(ctx, map[string]any{
			"plan":           string(planContent),
			"input":          formatInput(in.UserInput),
			"executed_steps": formatExecutedSteps(in.ExecutedSteps),
			"plan_tool":      planToolName,
			"respond_tool":   respondToolName,
		})
		if err != nil {
			return nil, err
		}

		return msgs, nil
	}
}

// NewReplanner creates a plan-execute-replan agent wired with plan and respond tools.
// It configures the provided ToolCallingChatModel with the tools and returns an Agent.
func NewReplanner(_ context.Context, cfg *ReplannerConfig) (adk.Agent, error) {
	planTool := cfg.PlanTool
	if planTool == nil {
		planTool = &PlanToolInfo
	}

	respondTool := cfg.RespondTool
	if respondTool == nil {
		respondTool = &RespondToolInfo
	}

	chatModel, err := cfg.ChatModel.WithTools([]*schema.ToolInfo{planTool, respondTool})
	if err != nil {
		return nil, err
	}

	planParser := cfg.NewPlan
	if planParser == nil {
		planParser = defaultNewPlan
	}

	return &replanner{
		chatModel:   chatModel,
		planTool:    planTool,
		respondTool: respondTool,
		genInputFn:  cfg.GenInputFn,
		newPlan:     planParser,
	}, nil
}

// Config provides configuration options for creating a plan-execute-replan agent.
type Config struct {
	// Planner specifies the agent that generates the plan.
	// You can use provided NewPlanner to create a planner agent.
	Planner adk.Agent

	// Executor specifies the agent that executes the plan generated by planner or replanner.
	// You can use provided NewExecutor to create an executor agent.
	Executor adk.Agent

	// Replanner specifies the agent that replans the plan.
	// You can use provided NewReplanner to create a replanner agent.
	Replanner adk.Agent

	// MaxIterations defines the maximum number of loops for 'execute-replan'.
	// Optional. If not provided, 10 will be used as the default.
	MaxIterations int
}

// New creates a new plan-execute-replan agent with the given configuration.
// The plan-execute-replan pattern works in three phases:
// 1. Planning: Generate a structured plan with clear, actionable steps
// 2. Execution: Execute the first step of the plan
// 3. Replanning: Evaluate progress and either complete the task or revise the plan
// This approach enables complex problem-solving through iterative refinement.
func New(ctx context.Context, cfg *Config) (adk.ResumableAgent, error) {
	maxIterations := cfg.MaxIterations
	if maxIterations <= 0 {
		maxIterations = 10
	}
	loop, err := adk.NewLoopAgent(ctx, &adk.LoopAgentConfig{
		Name:          "execute_replan",
		SubAgents:     []adk.Agent{cfg.Executor, cfg.Replanner},
		MaxIterations: maxIterations,
	})
	if err != nil {
		return nil, err
	}

	return adk.NewSequentialAgent(ctx, &adk.SequentialAgentConfig{
		Name:      "plan_execute_replan",
		SubAgents: []adk.Agent{cfg.Planner, loop},
	})
}