/*
 * 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 compose

import (
	"context"
	"errors"
	"fmt"
	"runtime/debug"
	"sync"

	"github.com/cloudwego/eino/callbacks"
	"github.com/cloudwego/eino/components"
	"github.com/cloudwego/eino/components/tool"
	"github.com/cloudwego/eino/internal/safe"
	"github.com/cloudwego/eino/schema"
)

type toolsNodeOptions struct {
	ToolOptions   []tool.Option
	ToolList      []tool.BaseTool
	executedTools map[string]string
}

// ToolsNodeOption is the option func type for ToolsNode.
type ToolsNodeOption func(o *toolsNodeOptions)

// WithToolOption adds tool options to the ToolsNode.
func WithToolOption(opts ...tool.Option) ToolsNodeOption {
	return func(o *toolsNodeOptions) {
		o.ToolOptions = append(o.ToolOptions, opts...)
	}
}

// WithToolList sets the tool list for the ToolsNode.
func WithToolList(tool ...tool.BaseTool) ToolsNodeOption {
	return func(o *toolsNodeOptions) {
		o.ToolList = tool
	}
}

func withExecutedTools(executedTools map[string]string) ToolsNodeOption {
	return func(o *toolsNodeOptions) {
		o.executedTools = executedTools
	}
}

// ToolsNode represents a node capable of executing tools within a graph.
// The Graph Node interface is defined as follows:
//
//	Invoke(ctx context.Context, input *schema.Message, opts ...ToolsNodeOption) ([]*schema.Message, error)
//	Stream(ctx context.Context, input *schema.Message, opts ...ToolsNodeOption) (*schema.StreamReader[[]*schema.Message], error)
//
// Input: An AssistantMessage containing ToolCalls
// Output: An array of ToolMessage where the order of elements corresponds to the order of ToolCalls in the input
type ToolsNode struct {
	tuple                *toolsTuple
	unknownToolHandler   func(ctx context.Context, name, input string) (string, error)
	executeSequentially  bool
	toolArgumentsHandler func(ctx context.Context, name, input string) (string, error)
}

// ToolsNodeConfig is the config for ToolsNode.
type ToolsNodeConfig struct {
	// Tools specify the list of tools can be called which are BaseTool but must implement InvokableTool or StreamableTool.
	Tools []tool.BaseTool

	// UnknownToolsHandler handles tool calls for non-existent tools when LLM hallucinates.
	// This field is optional. When not set, calling a non-existent tool will result in an error.
	// When provided, if the LLM attempts to call a tool that doesn't exist in the Tools list,
	// this handler will be invoked instead of returning an error, allowing graceful handling of hallucinated tools.
	// Parameters:
	//   - ctx: The context for the tool call
	//   - name: The name of the non-existent tool
	//   - input: The tool call input generated by llm
	// Returns:
	//   - string: The response to be returned as if the tool was executed
	//   - error: Any error that occurred during handling
	UnknownToolsHandler func(ctx context.Context, name, input string) (string, error)

	// ExecuteSequentially determines whether tool calls should be executed sequentially (in order) or in parallel.
	// When set to true, tool calls will be executed one after another in the order they appear in the input message.
	// When set to false (default), tool calls will be executed in parallel.
	ExecuteSequentially bool

	// ToolArgumentsHandler allows handling of tool arguments before execution.
	// When provided, this function will be called for each tool call to process the arguments.
	// Parameters:
	//   - ctx: The context for the tool call
	//   - name: The name of the tool being called
	//   - arguments: The original arguments string for the tool
	// Returns:
	//   - string: The processed arguments string to be used for tool execution
	//   - error: Any error that occurred during preprocessing
	ToolArgumentsHandler func(ctx context.Context, name, arguments string) (string, error)
}

// NewToolNode creates a new ToolsNode.
// e.g.
//
//	conf := &ToolsNodeConfig{
//		Tools: []tool.BaseTool{invokableTool1, streamableTool2},
//	}
//	toolsNode, err := NewToolNode(ctx, conf)
func NewToolNode(ctx context.Context, conf *ToolsNodeConfig) (*ToolsNode, error) {
	tuple, err := convTools(ctx, conf.Tools)
	if err != nil {
		return nil, err
	}

	return &ToolsNode{
		tuple:                tuple,
		unknownToolHandler:   conf.UnknownToolsHandler,
		executeSequentially:  conf.ExecuteSequentially,
		toolArgumentsHandler: conf.ToolArgumentsHandler,
	}, nil
}

type ToolsInterruptAndRerunExtra struct {
	ToolCalls     []schema.ToolCall
	ExecutedTools map[string]string
	RerunTools    []string
	RerunExtraMap map[string]any
}

func init() {
	schema.RegisterName[*ToolsInterruptAndRerunExtra]("_eino_compose_tools_interrupt_and_rerun_extra") // TODO: check if this is really needed when refactoring adk resume
}

type toolsTuple struct {
	indexes map[string]int
	meta    []*executorMeta
	rps     []*runnablePacker[string, string, tool.Option]
}

func convTools(ctx context.Context, tools []tool.BaseTool) (*toolsTuple, error) {
	ret := &toolsTuple{
		indexes: make(map[string]int),
		meta:    make([]*executorMeta, len(tools)),
		rps:     make([]*runnablePacker[string, string, tool.Option], len(tools)),
	}
	for idx, bt := range tools {
		tl, err := bt.Info(ctx)
		if err != nil {
			return nil, fmt.Errorf("(NewToolNode) failed to get tool info at idx= %d: %w", idx, err)
		}

		toolName := tl.Name
		var (
			st tool.StreamableTool
			it tool.InvokableTool

			invokable  func(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (string, error)
			streamable func(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (*schema.StreamReader[string], error)

			ok   bool
			meta *executorMeta
		)

		if st, ok = bt.(tool.StreamableTool); ok {
			streamable = st.StreamableRun
		}

		if it, ok = bt.(tool.InvokableTool); ok {
			invokable = it.InvokableRun
		}

		if st == nil && it == nil {
			return nil, fmt.Errorf("tool %s is not invokable or streamable", toolName)
		}

		if st != nil {
			meta = parseExecutorInfoFromComponent(components.ComponentOfTool, st)
		} else {
			meta = parseExecutorInfoFromComponent(components.ComponentOfTool, it)
		}

		ret.indexes[toolName] = idx
		ret.meta[idx] = meta
		ret.rps[idx] = newRunnablePacker(invokable, streamable,
			nil, nil, !meta.isComponentCallbackEnabled)
	}
	return ret, nil
}

type toolCallTask struct {
	// in
	r      *runnablePacker[string, string, tool.Option]
	meta   *executorMeta
	name   string
	arg    string
	callID string

	// out
	executed bool
	output   string
	sOutput  *schema.StreamReader[string]
	err      error
}

func (tn *ToolsNode) genToolCallTasks(ctx context.Context, tuple *toolsTuple,
	input *schema.Message, executedTools map[string]string, isStream bool) ([]toolCallTask, error) {

	if input.Role != schema.Assistant {
		return nil, fmt.Errorf("expected message role is Assistant, got %s", input.Role)
	}

	n := len(input.ToolCalls)
	if n == 0 {
		return nil, errors.New("no tool call found in input message")
	}

	toolCallTasks := make([]toolCallTask, n)

	for i := 0; i < n; i++ {
		toolCall := input.ToolCalls[i]
		if result, executed := executedTools[toolCall.ID]; executed {
			toolCallTasks[i].name = toolCall.Function.Name
			toolCallTasks[i].arg = toolCall.Function.Arguments
			toolCallTasks[i].callID = toolCall.ID
			toolCallTasks[i].executed = true
			if isStream {
				toolCallTasks[i].sOutput = schema.StreamReaderFromArray([]string{result})
			} else {
				toolCallTasks[i].output = result
			}
			continue
		}
		index, ok := tuple.indexes[toolCall.Function.Name]
		if !ok {
			if tn.unknownToolHandler == nil {
				return nil, fmt.Errorf("tool %s not found in toolsNode indexes", toolCall.Function.Name)
			}
			toolCallTasks[i] = newUnknownToolTask(toolCall.Function.Name, toolCall.Function.Arguments, toolCall.ID, tn.unknownToolHandler)
		} else {
			toolCallTasks[i].r = tuple.rps[index]
			toolCallTasks[i].meta = tuple.meta[index]
			toolCallTasks[i].name = toolCall.Function.Name
			toolCallTasks[i].callID = toolCall.ID
			if tn.toolArgumentsHandler != nil {
				arg, err := tn.toolArgumentsHandler(ctx, toolCall.Function.Name, toolCall.Function.Arguments)
				if err != nil {
					return nil, fmt.Errorf("failed to executed tool[name:%s arguments:%s] arguments handler: %w", toolCall.Function.Name, toolCall.Function.Arguments, err)
				}
				toolCallTasks[i].arg = arg
			} else {
				toolCallTasks[i].arg = toolCall.Function.Arguments
			}
		}
	}

	return toolCallTasks, nil
}

func newUnknownToolTask(name, arg, callID string, unknownToolHandler func(ctx context.Context, name, input string) (string, error)) toolCallTask {
	return toolCallTask{
		r: newRunnablePacker(func(ctx context.Context, input string, opts ...tool.Option) (output string, err error) {
			return unknownToolHandler(ctx, name, input)
		}, nil, nil, nil, false),
		meta: &executorMeta{
			component:                  components.ComponentOfTool,
			isComponentCallbackEnabled: false,
			componentImplType:          "UnknownTool",
		},
		name:   name,
		arg:    arg,
		callID: callID,
	}
}

func runToolCallTaskByInvoke(ctx context.Context, task *toolCallTask, opts ...tool.Option) {
	if task.executed {
		return
	}
	ctx = callbacks.ReuseHandlers(ctx, &callbacks.RunInfo{
		Name:      task.name,
		Type:      task.meta.componentImplType,
		Component: task.meta.component,
	})

	ctx = setToolCallInfo(ctx, &toolCallInfo{toolCallID: task.callID})
	task.output, task.err = task.r.Invoke(ctx, task.arg, opts...)
	if task.err == nil {
		task.executed = true
	}
}

func runToolCallTaskByStream(ctx context.Context, task *toolCallTask, opts ...tool.Option) {
	ctx = callbacks.ReuseHandlers(ctx, &callbacks.RunInfo{
		Name:      task.name,
		Type:      task.meta.componentImplType,
		Component: task.meta.component,
	})

	ctx = setToolCallInfo(ctx, &toolCallInfo{toolCallID: task.callID})
	task.sOutput, task.err = task.r.Stream(ctx, task.arg, opts...)
	if task.err == nil {
		task.executed = true
	}
}

func sequentialRunToolCall(ctx context.Context,
	run func(ctx2 context.Context, callTask *toolCallTask, opts ...tool.Option),
	tasks []toolCallTask, opts ...tool.Option) {

	for i := range tasks {
		if tasks[i].executed {
			continue
		}
		run(ctx, &tasks[i], opts...)
	}
}

func parallelRunToolCall(ctx context.Context,
	run func(ctx2 context.Context, callTask *toolCallTask, opts ...tool.Option),
	tasks []toolCallTask, opts ...tool.Option) {

	if len(tasks) == 1 {
		run(ctx, &tasks[0], opts...)
		return
	}

	var wg sync.WaitGroup
	for i := 1; i < len(tasks); i++ {
		if tasks[i].executed {
			continue
		}
		wg.Add(1)
		go func(ctx_ context.Context, t *toolCallTask, opts ...tool.Option) {
			defer wg.Done()
			defer func() {
				panicErr := recover()
				if panicErr != nil {
					t.err = safe.NewPanicErr(panicErr, debug.Stack())
				}
			}()
			run(ctx_, t, opts...)
		}(ctx, &tasks[i], opts...)
	}

	run(ctx, &tasks[0], opts...)
	wg.Wait()
}

// Invoke calls the tools and collects the results of invokable tools.
// it's parallel if there are multiple tool calls in the input message.
func (tn *ToolsNode) Invoke(ctx context.Context, input *schema.Message,
	opts ...ToolsNodeOption) ([]*schema.Message, error) {

	opt := getToolsNodeOptions(opts...)
	tuple := tn.tuple
	if opt.ToolList != nil {
		var err error
		tuple, err = convTools(ctx, opt.ToolList)
		if err != nil {
			return nil, fmt.Errorf("failed to convert tool list from call option: %w", err)
		}
	}

	tasks, err := tn.genToolCallTasks(ctx, tuple, input, opt.executedTools, false)
	if err != nil {
		return nil, err
	}

	if tn.executeSequentially {
		sequentialRunToolCall(ctx, runToolCallTaskByInvoke, tasks, opt.ToolOptions...)
	} else {
		parallelRunToolCall(ctx, runToolCallTaskByInvoke, tasks, opt.ToolOptions...)
	}

	n := len(tasks)
	output := make([]*schema.Message, n)

	rerunExtra := &ToolsInterruptAndRerunExtra{
		ToolCalls:     input.ToolCalls,
		ExecutedTools: make(map[string]string),
		RerunExtraMap: make(map[string]any),
	}
	rerun := false
	for i := 0; i < n; i++ {
		if tasks[i].err != nil {
			extra, ok := IsInterruptRerunError(tasks[i].err)
			if !ok {
				return nil, fmt.Errorf("failed to invoke tool[name:%s id:%s]: %w", tasks[i].name, tasks[i].callID, tasks[i].err)
			}
			rerun = true
			rerunExtra.RerunTools = append(rerunExtra.RerunTools, tasks[i].callID)
			rerunExtra.RerunExtraMap[tasks[i].callID] = extra
			continue
		}
		if tasks[i].executed {
			rerunExtra.ExecutedTools[tasks[i].callID] = tasks[i].output
		}
		if !rerun {
			output[i] = schema.ToolMessage(tasks[i].output, tasks[i].callID, schema.WithToolName(tasks[i].name))
		}
	}
	if rerun {
		return nil, NewInterruptAndRerunErr(rerunExtra)
	}

	return output, nil
}

// Stream calls the tools and collects the results of stream readers.
// it's parallel if there are multiple tool calls in the input message.
func (tn *ToolsNode) Stream(ctx context.Context, input *schema.Message,
	opts ...ToolsNodeOption) (*schema.StreamReader[[]*schema.Message], error) {

	opt := getToolsNodeOptions(opts...)
	tuple := tn.tuple
	if opt.ToolList != nil {
		var err error
		tuple, err = convTools(ctx, opt.ToolList)
		if err != nil {
			return nil, fmt.Errorf("failed to convert tool list from call option: %w", err)
		}
	}

	tasks, err := tn.genToolCallTasks(ctx, tuple, input, opt.executedTools, true)
	if err != nil {
		return nil, err
	}

	if tn.executeSequentially {
		sequentialRunToolCall(ctx, runToolCallTaskByStream, tasks, opt.ToolOptions...)
	} else {
		parallelRunToolCall(ctx, runToolCallTaskByStream, tasks, opt.ToolOptions...)
	}

	n := len(tasks)

	rerun := false
	rerunExtra := &ToolsInterruptAndRerunExtra{
		ToolCalls:     input.ToolCalls,
		RerunExtraMap: make(map[string]any),
		ExecutedTools: make(map[string]string),
	}

	// check rerun
	for i := 0; i < n; i++ {
		if tasks[i].err != nil {
			extra, ok := IsInterruptRerunError(tasks[i].err)
			if !ok {
				return nil, fmt.Errorf("failed to stream tool call %s: %w", tasks[i].callID, tasks[i].err)
			}
			rerun = true
			rerunExtra.RerunTools = append(rerunExtra.RerunTools, tasks[i].callID)
			rerunExtra.RerunExtraMap[tasks[i].callID] = extra
			continue
		}
	}

	if rerun {
		// concat and save tool output
		for _, t := range tasks {
			if t.executed {
				o, err_ := concatStreamReader(t.sOutput)
				if err_ != nil {
					return nil, fmt.Errorf("failed to concat tool[name:%s id:%s]'s stream output: %w", t.name, t.callID, err_)
				}
				rerunExtra.ExecutedTools[t.callID] = o
			}
		}
		return nil, NewInterruptAndRerunErr(rerunExtra)
	}

	// common return
	sOutput := make([]*schema.StreamReader[[]*schema.Message], n)
	for i := 0; i < n; i++ {
		index := i
		callID := tasks[i].callID
		callName := tasks[i].name
		cvt := func(s string) ([]*schema.Message, error) {
			ret := make([]*schema.Message, n)
			ret[index] = schema.ToolMessage(s, callID, schema.WithToolName(callName))

			return ret, nil
		}

		sOutput[i] = schema.StreamReaderWithConvert(tasks[i].sOutput, cvt)
	}
	return schema.MergeStreamReaders(sOutput), nil
}

func (tn *ToolsNode) GetType() string {
	return ""
}

func getToolsNodeOptions(opts ...ToolsNodeOption) *toolsNodeOptions {
	o := &toolsNodeOptions{
		ToolOptions: make([]tool.Option, 0),
	}
	for _, opt := range opts {
		opt(o)
	}
	return o
}

type toolCallInfoKey struct{}
type toolCallInfo struct {
	toolCallID string
}

func setToolCallInfo(ctx context.Context, toolCallInfo *toolCallInfo) context.Context {
	return context.WithValue(ctx, toolCallInfoKey{}, toolCallInfo)
}

// GetToolCallID gets the current tool call id from the context.
func GetToolCallID(ctx context.Context) string {
	v := ctx.Value(toolCallInfoKey{})
	if v == nil {
		return ""
	}

	info, ok := v.(*toolCallInfo)
	if !ok {
		return ""
	}

	return info.toolCallID
}
