package llms

import (
	"encoding/json"
	"fmt"
	"github.com/invopop/jsonschema"
	"reflect"
)

/**
Define
*/

type ToolType int

const (
	ToolTypeInvalid  ToolType = 0
	ToolTypeFunction ToolType = 1 // 工具类型为：函数
)

/**
Tool
*/

type Tool interface {
	Type() ToolType
	Name() string
	Desc() string
	Run(param string) (string, error)
}

type ToolBox map[string]Tool

func (tm ToolBox) Tools() []Tool {
	ret := make([]Tool, 0, len(tm))
	for _, v := range tm {
		ret = append(ret, v)
	}
	return ret
}

func (tm ToolBox) Get(k string) (Tool, bool) {
	t, ok := tm[k]
	return t, ok
}

func (tm ToolBox) Add(t Tool) {
	tm[t.Name()] = t
}

func (tm ToolBox) Remove(k string) {
	delete(tm, k)
}

/**
FunctionTool
*/

type ToolFunctionExample struct {
	Input          string // 用户问题
	FunctionParams string // Json 格式的函数参数示例
	IsCall         bool   // 是否会触发调用函数
}

type FunctionTool struct {
	name, desc, paramScheme string
	fn                      interface{}
	example                 []ToolFunctionExample
}

func (t *FunctionTool) Type() ToolType {
	return ToolTypeFunction
}

func (t *FunctionTool) Name() string {
	return t.name
}

func (t *FunctionTool) Desc() string {
	return t.desc
}

func (t *FunctionTool) ParamJsonScheme() (string, error) {
	if t.paramScheme == "" {
		funcType := reflect.TypeOf(t.fn)
		param := funcType.In(0)
		paramValue := reflect.New(param).Elem().Interface()

		var paramJson *jsonschema.Schema

		def := jsonschema.Reflect(paramValue).Definitions
		for k := range def {
			paramJson = def[k]
		}
		if paramJson != nil {
			scheme, err := json.Marshal(paramJson)
			if err != nil {
				return "", err
			}
			t.paramScheme = string(scheme)
		}
	}
	return t.paramScheme, nil
}

func (t *FunctionTool) Examples() []ToolFunctionExample {
	return t.example
}

func (t *FunctionTool) Run(param string) (string, error) {
	funcVal := reflect.ValueOf(t.fn)
	funcType := funcVal.Type()

	// 初始化参数
	reflectParam := funcType.In(0)
	paramValue := reflect.New(reflectParam).Elem()

	if err := json.Unmarshal([]byte(param), paramValue.Addr().Interface()); err != nil {
		return "", err
	}

	// 调用函数
	callParams := []reflect.Value{paramValue}
	results := funcVal.Call(callParams)

	// 获取返回结果
	if len(results) != 2 {
		return "", fmt.Errorf("incorrect number of return for fn (len(return) != 2)")
	}

	ret := results[0].Interface().(string)
	var err error
	if !results[1].IsNil() {
		err = results[1].Interface().(error)
	}
	return ret, err
}

func NewFuncTool(name, desc string, fn interface{}, example ...ToolFunctionExample) (*FunctionTool, error) {
	funcVal := reflect.ValueOf(fn)
	funcType := funcVal.Type()

	// fn 类型检测
	if reflect.Func != funcType.Kind() {
		return nil, fmt.Errorf("params fn is not a function")
	}

	// fn 输入检测
	if 1 != funcType.NumIn() {
		return nil, fmt.Errorf("incorrect number of arguments for fn (len(arguments) != 1)")
	}

	// fn 返回检测
	if 2 != funcType.NumOut() {
		return nil, fmt.Errorf("incorrect number of return for fn (len(return) != 2)")
	}
	firstReturnType, secondReturnType := funcType.Out(0), funcType.Out(1)

	if !(firstReturnType == reflect.TypeOf("") &&
		secondReturnType.Implements(reflect.TypeOf((*error)(nil)).Elem())) {
		// 检查第一个返回值的类型是否是 string && 检查第二个返回值的类型是否是 error
		return nil, fmt.Errorf("incorrect type of return for fn (first need string type, second need error type)")
	}

	ret := &FunctionTool{
		name:    name,
		desc:    desc,
		fn:      fn,
		example: example,
	}
	return ret, nil
}