/*
 * Copyright 2024 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 callbacks

import (
	"context"
	"fmt"
	"testing"

	"github.com/stretchr/testify/assert"

	"github.com/cloudwego/eino/adk"
	"github.com/cloudwego/eino/callbacks"
	"github.com/cloudwego/eino/components"
	"github.com/cloudwego/eino/components/document"
	"github.com/cloudwego/eino/components/embedding"
	"github.com/cloudwego/eino/components/indexer"
	"github.com/cloudwego/eino/components/model"
	"github.com/cloudwego/eino/components/prompt"
	"github.com/cloudwego/eino/components/retriever"
	"github.com/cloudwego/eino/components/tool"
	"github.com/cloudwego/eino/compose"
	"github.com/cloudwego/eino/schema"
)

func TestNewComponentTemplate(t *testing.T) {
	t.Run("TestNewComponentTemplate", func(t *testing.T) {
		cnt := 0
		tpl := NewHandlerHelper()
		tpl.ChatModel(&ModelCallbackHandler{
			OnStart: func(ctx context.Context, runInfo *callbacks.RunInfo, input *model.CallbackInput) context.Context {
				cnt++
				return ctx
			},
			OnEnd: func(ctx context.Context, runInfo *callbacks.RunInfo, output *model.CallbackOutput) context.Context {
				cnt++
				return ctx
			},
			OnEndWithStreamOutput: func(ctx context.Context, runInfo *callbacks.RunInfo, output *schema.StreamReader[*model.CallbackOutput]) context.Context {
				output.Close()
				cnt++
				return ctx
			},
			OnError: func(ctx context.Context, info *callbacks.RunInfo, err error) context.Context {
				cnt++
				return ctx
			}}).
			Embedding(&EmbeddingCallbackHandler{
				OnStart: func(ctx context.Context, runInfo *callbacks.RunInfo, input *embedding.CallbackInput) context.Context {
					cnt++
					return ctx
				},
				OnEnd: func(ctx context.Context, runInfo *callbacks.RunInfo, output *embedding.CallbackOutput) context.Context {
					cnt++
					return ctx
				},
				OnError: func(ctx context.Context, info *callbacks.RunInfo, err error) context.Context {
					cnt++
					return ctx
				},
			}).
			Prompt(&PromptCallbackHandler{
				OnStart: func(ctx context.Context, runInfo *callbacks.RunInfo, input *prompt.CallbackInput) context.Context {
					cnt++
					return ctx
				},
				OnEnd: func(ctx context.Context, runInfo *callbacks.RunInfo, output *prompt.CallbackOutput) context.Context {
					cnt++
					return ctx
				},
				OnError: func(ctx context.Context, info *callbacks.RunInfo, err error) context.Context {
					cnt++
					return ctx
				},
			}).
			Retriever(&RetrieverCallbackHandler{
				OnStart: func(ctx context.Context, runInfo *callbacks.RunInfo, input *retriever.CallbackInput) context.Context {
					cnt++
					return ctx
				},
				OnEnd: func(ctx context.Context, runInfo *callbacks.RunInfo, output *retriever.CallbackOutput) context.Context {
					cnt++
					return ctx
				},
				OnError: func(ctx context.Context, info *callbacks.RunInfo, err error) context.Context {
					cnt++
					return ctx
				},
			}).
			Tool(&ToolCallbackHandler{
				OnStart: func(ctx context.Context, runInfo *callbacks.RunInfo, input *tool.CallbackInput) context.Context {
					cnt++
					return ctx
				},
				OnEnd: func(ctx context.Context, runInfo *callbacks.RunInfo, output *tool.CallbackOutput) context.Context {
					cnt++
					return ctx
				},
				OnEndWithStreamOutput: func(ctx context.Context, runInfo *callbacks.RunInfo, output *schema.StreamReader[*tool.CallbackOutput]) context.Context {
					cnt++
					return ctx
				},
				OnError: func(ctx context.Context, info *callbacks.RunInfo, err error) context.Context {
					cnt++
					return ctx
				},
			}).
			Lambda(callbacks.NewHandlerBuilder().
				OnStartFn(func(ctx context.Context, info *callbacks.RunInfo, input callbacks.CallbackInput) context.Context {
					cnt++
					return ctx
				}).
				OnStartWithStreamInputFn(func(ctx context.Context, info *callbacks.RunInfo, input *schema.StreamReader[callbacks.CallbackInput]) context.Context {
					input.Close()
					cnt++
					return ctx
				}).
				OnEndFn(func(ctx context.Context, info *callbacks.RunInfo, output callbacks.CallbackOutput) context.Context {
					cnt++
					return ctx
				}).
				OnEndWithStreamOutputFn(func(ctx context.Context, info *callbacks.RunInfo, output *schema.StreamReader[callbacks.CallbackOutput]) context.Context {
					output.Close()
					cnt++
					return ctx
				}).
				OnErrorFn(func(ctx context.Context, info *callbacks.RunInfo, err error) context.Context {
					cnt++
					return ctx
				}).Build()).
			AgenticModel(&AgenticModelCallbackHandler{
				OnStart: func(ctx context.Context, runInfo *callbacks.RunInfo, input *model.AgenticCallbackInput) context.Context {
					cnt++
					return ctx
				},
				OnEnd: func(ctx context.Context, runInfo *callbacks.RunInfo, output *model.AgenticCallbackOutput) context.Context {
					cnt++
					return ctx
				},
				OnEndWithStreamOutput: func(ctx context.Context, runInfo *callbacks.RunInfo, output *schema.StreamReader[*model.AgenticCallbackOutput]) context.Context {
					output.Close()
					cnt++
					return ctx
				},
				OnError: func(ctx context.Context, runInfo *callbacks.RunInfo, err error) context.Context {
					cnt++
					return ctx
				},
			}).
			AgenticPrompt(&AgenticPromptCallbackHandler{
				OnStart: func(ctx context.Context, runInfo *callbacks.RunInfo, input *prompt.CallbackInput) context.Context {
					cnt++
					return ctx
				},
				OnEnd: func(ctx context.Context, runInfo *callbacks.RunInfo, output *prompt.CallbackOutput) context.Context {
					cnt++
					return ctx
				},
				OnError: func(ctx context.Context, runInfo *callbacks.RunInfo, err error) context.Context {
					cnt++
					return ctx
				},
			}).
			AgenticToolsNode(&AgenticToolsNodeCallbackHandlers{
				OnStart: func(ctx context.Context, info *callbacks.RunInfo, input *schema.AgenticMessage) context.Context {
					cnt++
					return ctx
				},
				OnEnd: func(ctx context.Context, info *callbacks.RunInfo, input []*schema.AgenticMessage) context.Context {
					cnt++
					return ctx
				},
				OnEndWithStreamOutput: func(ctx context.Context, info *callbacks.RunInfo, output *schema.StreamReader[[]*schema.AgenticMessage]) context.Context {
					output.Close()
					cnt++
					return ctx
				},
				OnError: func(ctx context.Context, info *callbacks.RunInfo, err error) context.Context {
					cnt++
					return ctx
				},
			}).
			Handler()

		types := []components.Component{
			components.ComponentOfPrompt,
			components.ComponentOfChatModel,
			components.ComponentOfEmbedding,
			components.ComponentOfRetriever,
			components.ComponentOfTool,
			compose.ComponentOfLambda,
			components.ComponentOfAgenticModel,
			components.ComponentOfAgenticPrompt,
			compose.ComponentOfAgenticToolsNode,
		}

		handler := tpl.Handler()
		ctx := context.Background()
		for _, typ := range types {
			handler.OnStart(ctx, &callbacks.RunInfo{Component: typ}, nil)
			handler.OnEnd(ctx, &callbacks.RunInfo{Component: typ}, nil)
			handler.OnError(ctx, &callbacks.RunInfo{Component: typ}, fmt.Errorf("mock err"))

			sir, siw := schema.Pipe[callbacks.CallbackInput](1)
			siw.Close()
			handler.OnStartWithStreamInput(ctx, &callbacks.RunInfo{Component: typ}, sir)

			sor, sow := schema.Pipe[callbacks.CallbackOutput](1)
			sow.Close()
			handler.OnEndWithStreamOutput(ctx, &callbacks.RunInfo{Component: typ}, sor)
		}

		assert.Equal(t, 33, cnt)

		ctx = context.Background()
		ctx = callbacks.InitCallbacks(ctx, &callbacks.RunInfo{Component: components.ComponentOfTransformer}, handler)
		callbacks.OnStart[any](ctx, nil)
		assert.Equal(t, 33, cnt)

		ctx = callbacks.ReuseHandlers(ctx, &callbacks.RunInfo{Component: components.ComponentOfPrompt})
		ctx = callbacks.OnStart[any](ctx, nil)
		assert.Equal(t, 34, cnt)

		ctx = callbacks.ReuseHandlers(ctx, &callbacks.RunInfo{Component: components.ComponentOfIndexer})
		callbacks.OnEnd[any](ctx, nil)
		assert.Equal(t, 34, cnt)

		ctx = callbacks.ReuseHandlers(ctx, &callbacks.RunInfo{Component: components.ComponentOfEmbedding})
		callbacks.OnError(ctx, nil)
		assert.Equal(t, 35, cnt)

		ctx = callbacks.ReuseHandlers(ctx, &callbacks.RunInfo{Component: components.ComponentOfLoader})
		callbacks.OnStart[any](ctx, nil)
		assert.Equal(t, 35, cnt)

		tpl.Transformer(&TransformerCallbackHandler{
			OnStart: func(ctx context.Context, runInfo *callbacks.RunInfo, input *document.TransformerCallbackInput) context.Context {
				cnt++
				return ctx
			},
			OnEnd: func(ctx context.Context, runInfo *callbacks.RunInfo, output *document.TransformerCallbackOutput) context.Context {
				cnt++
				return ctx
			},
			OnError: func(ctx context.Context, info *callbacks.RunInfo, err error) context.Context {
				cnt++
				return ctx
			},
		}).Indexer(&IndexerCallbackHandler{
			OnStart: func(ctx context.Context, runInfo *callbacks.RunInfo, input *indexer.CallbackInput) context.Context {
				cnt++
				return ctx
			},
			OnEnd: func(ctx context.Context, runInfo *callbacks.RunInfo, output *indexer.CallbackOutput) context.Context {
				cnt++
				return ctx
			},
			OnError: func(ctx context.Context, info *callbacks.RunInfo, err error) context.Context {
				cnt++
				return ctx
			},
		}).Loader(&LoaderCallbackHandler{
			OnStart: func(ctx context.Context, runInfo *callbacks.RunInfo, input *document.LoaderCallbackInput) context.Context {
				cnt++
				return ctx
			},
			OnEnd: func(ctx context.Context, runInfo *callbacks.RunInfo, output *document.LoaderCallbackOutput) context.Context {
				cnt++
				return ctx
			},
			OnError: func(ctx context.Context, info *callbacks.RunInfo, err error) context.Context {
				cnt++
				return ctx
			},
		}).ToolsNode(&ToolsNodeCallbackHandlers{
			OnStart: func(ctx context.Context, runInfo *callbacks.RunInfo, input *schema.Message) context.Context {
				cnt++
				return ctx
			},
			OnEndWithStreamOutput: func(ctx context.Context, runInfo *callbacks.RunInfo, output *schema.StreamReader[[]*schema.Message]) context.Context {
				cnt++

				if output == nil {
					return ctx
				}

				for {
					_, err := output.Recv()
					if err != nil {
						return ctx
					}
				}
			},
		}).AgenticPrompt(&AgenticPromptCallbackHandler{
			OnStart: func(ctx context.Context, runInfo *callbacks.RunInfo, input *prompt.CallbackInput) context.Context {
				cnt++
				return ctx
			},
			OnEnd: func(ctx context.Context, runInfo *callbacks.RunInfo, output *prompt.CallbackOutput) context.Context {
				cnt++
				return ctx
			},
			OnError: func(ctx context.Context, runInfo *callbacks.RunInfo, err error) context.Context {
				cnt++
				return ctx
			},
		}).AgenticToolsNode(&AgenticToolsNodeCallbackHandlers{
			OnStart: func(ctx context.Context, info *callbacks.RunInfo, input *schema.AgenticMessage) context.Context {
				cnt++
				return ctx
			},
			OnEnd: func(ctx context.Context, info *callbacks.RunInfo, input []*schema.AgenticMessage) context.Context {
				cnt++
				return ctx
			},
			OnEndWithStreamOutput: func(ctx context.Context, info *callbacks.RunInfo, output *schema.StreamReader[[]*schema.AgenticMessage]) context.Context {
				output.Close()
				cnt++
				return ctx
			},
			OnError: func(ctx context.Context, info *callbacks.RunInfo, err error) context.Context {
				cnt++
				return ctx
			},
		})

		handler = tpl.Handler()
		ctx = context.Background()
		ctx = callbacks.InitCallbacks(ctx, &callbacks.RunInfo{Component: components.ComponentOfTransformer}, handler)

		ctx = callbacks.OnStart[any](ctx, nil)
		assert.Equal(t, 36, cnt)

		callbacks.OnEnd[any](ctx, nil)
		assert.Equal(t, 37, cnt)

		ctx = callbacks.ReuseHandlers(ctx, &callbacks.RunInfo{Component: components.ComponentOfLoader})
		callbacks.OnEnd[any](ctx, nil)
		assert.Equal(t, 38, cnt)

		ctx = callbacks.ReuseHandlers(ctx, &callbacks.RunInfo{Component: compose.ComponentOfToolsNode})
		callbacks.OnStart[any](ctx, nil)
		assert.Equal(t, 39, cnt)

		sr, sw := schema.Pipe[any](0)
		sw.Close()
		callbacks.OnEndWithStreamOutput[any](ctx, sr)
		assert.Equal(t, 40, cnt)

		sr1, sw1 := schema.Pipe[[]*schema.Message](1)
		sw1.Send([]*schema.Message{{}}, nil)
		sw1.Close()
		callbacks.OnEndWithStreamOutput[[]*schema.Message](ctx, sr1)
		// Check AgenticModel stream
		sir2, siw2 := schema.Pipe[callbacks.CallbackOutput](1)
		siw2.Close()
		handler.OnEndWithStreamOutput(ctx, &callbacks.RunInfo{Component: components.ComponentOfAgenticModel}, sir2)
		assert.Equal(t, 42, cnt)

		// Check AgenticToolsNode stream
		sir3, siw3 := schema.Pipe[callbacks.CallbackOutput](1)
		siw3.Close()
		handler.OnEndWithStreamOutput(ctx, &callbacks.RunInfo{Component: compose.ComponentOfAgenticToolsNode}, sir3)
		assert.Equal(t, 43, cnt)

		ctx = callbacks.ReuseHandlers(ctx, nil)
		callbacks.OnStart[any](ctx, nil)
		assert.Equal(t, 43, cnt)
	})

	t.Run("EdgeCases", func(t *testing.T) {
		ctx := context.Background()
		cnt := 0

		// 1. Test Graph and Chain Setters and Execution
		tpl := NewHandlerHelper().
			Graph(callbacks.NewHandlerBuilder().
				OnStartFn(func(ctx context.Context, info *callbacks.RunInfo, input callbacks.CallbackInput) context.Context {
					cnt++
					return ctx
				}).Build()).
			Chain(callbacks.NewHandlerBuilder().
				OnEndFn(func(ctx context.Context, info *callbacks.RunInfo, output callbacks.CallbackOutput) context.Context {
					cnt++
					return ctx
				}).Build())

		h := tpl.Handler()

		// Trigger Graph OnStart
		h.OnStart(ctx, &callbacks.RunInfo{Component: compose.ComponentOfGraph}, nil)
		assert.Equal(t, 1, cnt)

		// Trigger Chain OnEnd
		h.OnEnd(ctx, &callbacks.RunInfo{Component: compose.ComponentOfChain}, nil)
		assert.Equal(t, 2, cnt)

		// 2. Test Needed logic for Graph/Chain when handler is present/absent
		// Graph is present (OnStart)
		needed := h.(callbacks.TimingChecker).Needed(ctx, &callbacks.RunInfo{Component: compose.ComponentOfGraph}, callbacks.TimingOnStart)
		assert.True(t, needed)

		// Chain is present (OnEnd) - but we check OnStart which is not defined in the builder above?
		// NewHandlerBuilder returns a handler that usually returns true for Needed if the specific func is not nil.
		// Let's verify Chain OnStart is NOT needed because we only set OnEndFn.
		needed = h.(callbacks.TimingChecker).Needed(ctx, &callbacks.RunInfo{Component: compose.ComponentOfChain}, callbacks.TimingOnStart)
		assert.False(t, needed) // Should be false because OnStartFn wasn't set for Chain

		// Lambda is NOT present
		needed = h.(callbacks.TimingChecker).Needed(ctx, &callbacks.RunInfo{Component: compose.ComponentOfLambda}, callbacks.TimingOnStart)
		assert.False(t, needed)

		// 3. Test Conversion Fallbacks (Default cases)
		// We need a handler with ToolsNode and AgenticToolsNode to test their conversion fallbacks
		tpl2 := NewHandlerHelper().
			ToolsNode(&ToolsNodeCallbackHandlers{
				OnStart: func(ctx context.Context, info *callbacks.RunInfo, input *schema.Message) context.Context {
					if input == nil {
						cnt++
					}
					return ctx
				},
				OnEnd: func(ctx context.Context, info *callbacks.RunInfo, input []*schema.Message) context.Context {
					if input == nil {
						cnt++
					}
					return ctx
				},
			}).
			AgenticToolsNode(&AgenticToolsNodeCallbackHandlers{
				OnStart: func(ctx context.Context, info *callbacks.RunInfo, input *schema.AgenticMessage) context.Context {
					if input == nil {
						cnt++
					}
					return ctx
				},
				OnEnd: func(ctx context.Context, info *callbacks.RunInfo, input []*schema.AgenticMessage) context.Context {
					if input == nil {
						cnt++
					}
					return ctx
				},
			})

		h2 := tpl2.Handler()

		// Pass wrong type (string) to trigger default case in convToolsNodeCallbackInput -> returns nil
		h2.OnStart(ctx, &callbacks.RunInfo{Component: compose.ComponentOfToolsNode}, "wrong-input-type")
		assert.Equal(t, 3, cnt) // +1

		// Pass wrong type to trigger default case in convToolsNodeCallbackOutput -> returns nil
		h2.OnEnd(ctx, &callbacks.RunInfo{Component: compose.ComponentOfToolsNode}, "wrong-output-type")
		assert.Equal(t, 4, cnt) // +1

		// Pass wrong type to trigger default case in convAgenticToolsNodeCallbackInput -> returns nil
		h2.OnStart(ctx, &callbacks.RunInfo{Component: compose.ComponentOfAgenticToolsNode}, "wrong-input-type")
		assert.Equal(t, 5, cnt) // +1

		// Pass wrong type to trigger default case in convAgenticToolsNodeCallbackOutput -> returns nil
		h2.OnEnd(ctx, &callbacks.RunInfo{Component: compose.ComponentOfAgenticToolsNode}, "wrong-output-type")
		assert.Equal(t, 6, cnt) // +1

		// 4. Test Needed for Agentic components when handlers are Set vs Unset
		// tpl2 has AgenticToolsNode set
		needed = h2.(callbacks.TimingChecker).Needed(ctx, &callbacks.RunInfo{Component: compose.ComponentOfAgenticToolsNode}, callbacks.TimingOnStart)
		assert.True(t, needed)

		// tpl2 does NOT have AgenticModel set
		needed = h2.(callbacks.TimingChecker).Needed(ctx, &callbacks.RunInfo{Component: components.ComponentOfAgenticModel}, callbacks.TimingOnStart)
		assert.False(t, needed)

		// Set it now
		tpl2.AgenticModel(&AgenticModelCallbackHandler{
			OnStart: func(ctx context.Context, runInfo *callbacks.RunInfo, input *model.AgenticCallbackInput) context.Context {
				return ctx
			},
		})

		needed = h2.(callbacks.TimingChecker).Needed(ctx, &callbacks.RunInfo{Component: components.ComponentOfAgenticModel}, callbacks.TimingOnStart)
		assert.True(t, needed)

		// Check invalid component
		needed = h2.(callbacks.TimingChecker).Needed(ctx, &callbacks.RunInfo{Component: "UnknownComponent"}, callbacks.TimingOnStart)
		assert.False(t, needed)

		// Check RunInfo nil
		needed = h2.(callbacks.TimingChecker).Needed(ctx, nil, callbacks.TimingOnStart)
		assert.False(t, needed)

		// 5. Test Needed for Transformer, Loader, Indexer, etc to ensure switch coverage
		tpl3 := NewHandlerHelper().
			Transformer(&TransformerCallbackHandler{OnStart: func(ctx context.Context, info *callbacks.RunInfo, input *document.TransformerCallbackInput) context.Context {
				return ctx
			}}).
			Loader(&LoaderCallbackHandler{OnStart: func(ctx context.Context, info *callbacks.RunInfo, input *document.LoaderCallbackInput) context.Context {
				return ctx
			}}).
			Indexer(&IndexerCallbackHandler{OnStart: func(ctx context.Context, info *callbacks.RunInfo, input *indexer.CallbackInput) context.Context {
				return ctx
			}}).
			Retriever(&RetrieverCallbackHandler{OnStart: func(ctx context.Context, info *callbacks.RunInfo, input *retriever.CallbackInput) context.Context {
				return ctx
			}}).
			Embedding(&EmbeddingCallbackHandler{OnStart: func(ctx context.Context, info *callbacks.RunInfo, input *embedding.CallbackInput) context.Context {
				return ctx
			}}).
			Tool(&ToolCallbackHandler{OnStart: func(ctx context.Context, info *callbacks.RunInfo, input *tool.CallbackInput) context.Context {
				return ctx
			}})

		h3 := tpl3.Handler()
		checker := h3.(callbacks.TimingChecker)

		assert.True(t, checker.Needed(ctx, &callbacks.RunInfo{Component: components.ComponentOfTransformer}, callbacks.TimingOnStart))
		assert.True(t, checker.Needed(ctx, &callbacks.RunInfo{Component: components.ComponentOfLoader}, callbacks.TimingOnStart))
		assert.True(t, checker.Needed(ctx, &callbacks.RunInfo{Component: components.ComponentOfIndexer}, callbacks.TimingOnStart))
		assert.True(t, checker.Needed(ctx, &callbacks.RunInfo{Component: components.ComponentOfRetriever}, callbacks.TimingOnStart))
		assert.True(t, checker.Needed(ctx, &callbacks.RunInfo{Component: components.ComponentOfEmbedding}, callbacks.TimingOnStart))
		assert.True(t, checker.Needed(ctx, &callbacks.RunInfo{Component: components.ComponentOfTool}, callbacks.TimingOnStart))

		// Verify False paths (by using a helper without them)
		emptyH := NewHandlerHelper().Handler().(callbacks.TimingChecker)
		assert.False(t, emptyH.Needed(ctx, &callbacks.RunInfo{Component: components.ComponentOfTransformer}, callbacks.TimingOnStart))
		assert.False(t, emptyH.Needed(ctx, &callbacks.RunInfo{Component: components.ComponentOfLoader}, callbacks.TimingOnStart))
		assert.False(t, emptyH.Needed(ctx, &callbacks.RunInfo{Component: components.ComponentOfIndexer}, callbacks.TimingOnStart))
		assert.False(t, emptyH.Needed(ctx, &callbacks.RunInfo{Component: components.ComponentOfRetriever}, callbacks.TimingOnStart))
		assert.False(t, emptyH.Needed(ctx, &callbacks.RunInfo{Component: components.ComponentOfEmbedding}, callbacks.TimingOnStart))
		assert.False(t, emptyH.Needed(ctx, &callbacks.RunInfo{Component: components.ComponentOfTool}, callbacks.TimingOnStart))

		// 6. Test Needed for remaining components (ChatModel, Prompt, AgenticPrompt)
		tpl4 := NewHandlerHelper().
			ChatModel(&ModelCallbackHandler{OnStart: func(ctx context.Context, runInfo *callbacks.RunInfo, input *model.CallbackInput) context.Context {
				return ctx
			}}).
			Prompt(&PromptCallbackHandler{OnStart: func(ctx context.Context, runInfo *callbacks.RunInfo, input *prompt.CallbackInput) context.Context {
				return ctx
			}}).
			AgenticPrompt(&AgenticPromptCallbackHandler{OnStart: func(ctx context.Context, runInfo *callbacks.RunInfo, input *prompt.CallbackInput) context.Context {
				return ctx
			}})

		h4 := tpl4.Handler()
		checker4 := h4.(callbacks.TimingChecker)

		assert.True(t, checker4.Needed(ctx, &callbacks.RunInfo{Component: components.ComponentOfChatModel}, callbacks.TimingOnStart))
		assert.True(t, checker4.Needed(ctx, &callbacks.RunInfo{Component: components.ComponentOfPrompt}, callbacks.TimingOnStart))
		assert.True(t, checker4.Needed(ctx, &callbacks.RunInfo{Component: components.ComponentOfAgenticPrompt}, callbacks.TimingOnStart))
	})
}

func TestAgentCallbackHandler(t *testing.T) {
	t.Run("Needed returns correct values", func(t *testing.T) {
		handler := &AgentCallbackHandler{
			OnStart: func(ctx context.Context, info *callbacks.RunInfo, input *adk.AgentCallbackInput) context.Context {
				return ctx
			},
		}

		ctx := context.Background()
		info := &callbacks.RunInfo{Component: adk.ComponentOfAgent}

		assert.True(t, handler.Needed(ctx, info, callbacks.TimingOnStart))
		assert.False(t, handler.Needed(ctx, info, callbacks.TimingOnEnd))
	})

	t.Run("Needed with OnEnd set", func(t *testing.T) {
		handler := &AgentCallbackHandler{
			OnEnd: func(ctx context.Context, info *callbacks.RunInfo, output *adk.AgentCallbackOutput) context.Context {
				return ctx
			},
		}

		ctx := context.Background()
		info := &callbacks.RunInfo{Component: adk.ComponentOfAgent}

		assert.False(t, handler.Needed(ctx, info, callbacks.TimingOnStart))
		assert.True(t, handler.Needed(ctx, info, callbacks.TimingOnEnd))
	})

	t.Run("Needed with nil handlers", func(t *testing.T) {
		handler := &AgentCallbackHandler{}

		ctx := context.Background()
		info := &callbacks.RunInfo{Component: adk.ComponentOfAgent}

		assert.False(t, handler.Needed(ctx, info, callbacks.TimingOnStart))
		assert.False(t, handler.Needed(ctx, info, callbacks.TimingOnEnd))
	})
}

func TestHandlerHelperWithAgent(t *testing.T) {
	t.Run("Agent method sets handler correctly", func(t *testing.T) {
		cnt := 0
		tpl := NewHandlerHelper()
		tpl.Agent(&AgentCallbackHandler{
			OnStart: func(ctx context.Context, info *callbacks.RunInfo, input *adk.AgentCallbackInput) context.Context {
				cnt++
				return ctx
			},
			OnEnd: func(ctx context.Context, info *callbacks.RunInfo, output *adk.AgentCallbackOutput) context.Context {
				cnt++
				return ctx
			},
		})

		handler := tpl.Handler()
		ctx := context.Background()
		ctx = callbacks.InitCallbacks(ctx, &callbacks.RunInfo{Component: adk.ComponentOfAgent}, handler)

		ctx = callbacks.OnStart[any](ctx, nil)
		assert.Equal(t, 1, cnt)

		callbacks.OnEnd[any](ctx, nil)
		assert.Equal(t, 2, cnt)
	})
}

func TestHandlerTemplateWithAgentComponent(t *testing.T) {
	t.Run("OnStart routes to agent handler", func(t *testing.T) {
		called := false
		tpl := NewHandlerHelper()
		tpl.Agent(&AgentCallbackHandler{
			OnStart: func(ctx context.Context, info *callbacks.RunInfo, input *adk.AgentCallbackInput) context.Context {
				called = true
				return ctx
			},
		})

		handler := tpl.Handler()
		ctx := context.Background()
		info := &callbacks.RunInfo{Component: adk.ComponentOfAgent, Name: "TestAgent"}

		handler.OnStart(ctx, info, &adk.AgentCallbackInput{})
		assert.True(t, called)
	})

	t.Run("OnEnd routes to agent handler", func(t *testing.T) {
		called := false
		tpl := NewHandlerHelper()
		tpl.Agent(&AgentCallbackHandler{
			OnEnd: func(ctx context.Context, info *callbacks.RunInfo, output *adk.AgentCallbackOutput) context.Context {
				called = true
				return ctx
			},
		})

		handler := tpl.Handler()
		ctx := context.Background()
		info := &callbacks.RunInfo{Component: adk.ComponentOfAgent, Name: "TestAgent"}

		handler.OnEnd(ctx, info, &adk.AgentCallbackOutput{})
		assert.True(t, called)
	})

	t.Run("Needed returns true for agent component", func(t *testing.T) {
		tpl := NewHandlerHelper()
		tpl.Agent(&AgentCallbackHandler{
			OnStart: func(ctx context.Context, info *callbacks.RunInfo, input *adk.AgentCallbackInput) context.Context {
				return ctx
			},
		})

		handler := tpl.Handler()
		ctx := context.Background()
		info := &callbacks.RunInfo{Component: adk.ComponentOfAgent}

		checker, ok := handler.(callbacks.TimingChecker)
		assert.True(t, ok, "handler should implement TimingChecker")
		assert.True(t, checker.Needed(ctx, info, callbacks.TimingOnStart))
	})
}

func TestAgenticAgentCallbackHandler(t *testing.T) {
	t.Run("Needed returns correct values", func(t *testing.T) {
		handler := &AgenticAgentCallbackHandler{
			OnStart: func(ctx context.Context, info *callbacks.RunInfo, input *adk.TypedAgentCallbackInput[*schema.AgenticMessage]) context.Context {
				return ctx
			},
		}

		ctx := context.Background()
		info := &callbacks.RunInfo{Component: adk.ComponentOfAgenticAgent}

		assert.True(t, handler.Needed(ctx, info, callbacks.TimingOnStart))
		assert.False(t, handler.Needed(ctx, info, callbacks.TimingOnEnd))
	})

	t.Run("Needed with OnEnd set", func(t *testing.T) {
		handler := &AgenticAgentCallbackHandler{
			OnEnd: func(ctx context.Context, info *callbacks.RunInfo, output *adk.TypedAgentCallbackOutput[*schema.AgenticMessage]) context.Context {
				return ctx
			},
		}

		ctx := context.Background()
		info := &callbacks.RunInfo{Component: adk.ComponentOfAgenticAgent}

		assert.False(t, handler.Needed(ctx, info, callbacks.TimingOnStart))
		assert.True(t, handler.Needed(ctx, info, callbacks.TimingOnEnd))
	})

	t.Run("Needed with nil handlers", func(t *testing.T) {
		handler := &AgenticAgentCallbackHandler{}

		ctx := context.Background()
		info := &callbacks.RunInfo{Component: adk.ComponentOfAgenticAgent}

		assert.False(t, handler.Needed(ctx, info, callbacks.TimingOnStart))
		assert.False(t, handler.Needed(ctx, info, callbacks.TimingOnEnd))
	})
}

func TestHandlerHelperWithAgenticAgent(t *testing.T) {
	t.Run("AgenticAgent method sets handler correctly", func(t *testing.T) {
		cnt := 0
		tpl := NewHandlerHelper()
		tpl.AgenticAgent(&AgenticAgentCallbackHandler{
			OnStart: func(ctx context.Context, info *callbacks.RunInfo, input *adk.TypedAgentCallbackInput[*schema.AgenticMessage]) context.Context {
				cnt++
				return ctx
			},
			OnEnd: func(ctx context.Context, info *callbacks.RunInfo, output *adk.TypedAgentCallbackOutput[*schema.AgenticMessage]) context.Context {
				cnt++
				return ctx
			},
		})

		handler := tpl.Handler()
		ctx := context.Background()
		ctx = callbacks.InitCallbacks(ctx, &callbacks.RunInfo{Component: adk.ComponentOfAgenticAgent}, handler)

		ctx = callbacks.OnStart[any](ctx, nil)
		assert.Equal(t, 1, cnt)

		callbacks.OnEnd[any](ctx, nil)
		assert.Equal(t, 2, cnt)
	})
}

func TestHandlerTemplateWithAgenticAgentComponent(t *testing.T) {
	t.Run("OnStart routes to agentic agent handler", func(t *testing.T) {
		called := false
		tpl := NewHandlerHelper()
		tpl.AgenticAgent(&AgenticAgentCallbackHandler{
			OnStart: func(ctx context.Context, info *callbacks.RunInfo, input *adk.TypedAgentCallbackInput[*schema.AgenticMessage]) context.Context {
				called = true
				return ctx
			},
		})

		handler := tpl.Handler()
		ctx := context.Background()
		info := &callbacks.RunInfo{Component: adk.ComponentOfAgenticAgent, Name: "TestAgenticAgent"}

		handler.OnStart(ctx, info, &adk.TypedAgentCallbackInput[*schema.AgenticMessage]{})
		assert.True(t, called)
	})

	t.Run("OnEnd routes to agentic agent handler", func(t *testing.T) {
		called := false
		tpl := NewHandlerHelper()
		tpl.AgenticAgent(&AgenticAgentCallbackHandler{
			OnEnd: func(ctx context.Context, info *callbacks.RunInfo, output *adk.TypedAgentCallbackOutput[*schema.AgenticMessage]) context.Context {
				called = true
				return ctx
			},
		})

		handler := tpl.Handler()
		ctx := context.Background()
		info := &callbacks.RunInfo{Component: adk.ComponentOfAgenticAgent, Name: "TestAgenticAgent"}

		handler.OnEnd(ctx, info, &adk.TypedAgentCallbackOutput[*schema.AgenticMessage]{})
		assert.True(t, called)
	})

	t.Run("Needed returns true for agentic agent component", func(t *testing.T) {
		tpl := NewHandlerHelper()
		tpl.AgenticAgent(&AgenticAgentCallbackHandler{
			OnStart: func(ctx context.Context, info *callbacks.RunInfo, input *adk.TypedAgentCallbackInput[*schema.AgenticMessage]) context.Context {
				return ctx
			},
		})

		handler := tpl.Handler()
		ctx := context.Background()
		info := &callbacks.RunInfo{Component: adk.ComponentOfAgenticAgent}

		checker, ok := handler.(callbacks.TimingChecker)
		assert.True(t, ok, "handler should implement TimingChecker")
		assert.True(t, checker.Needed(ctx, info, callbacks.TimingOnStart))
	})
}