package js

import (
	"apigo.cc/ai/ai/goja"
	"apigo.cc/ai/ai/interface/llm"
	"github.com/ssgo/u"
	"reflect"
	"strings"
)

func RequireAI() map[string]any {
	aiObj := map[string]any{
		"similarity": func(argsIn goja.FunctionCall, vm *goja.Runtime) goja.Value {
			args := MakeArgs(&argsIn, vm).Check(2)
			return vm.ToValue(llm.Similarity(args.Bytes(0), args.Bytes(1)))
		},
	}
	for name, lm := range llm.List() {
		aiObj[name] = RequireLLM(lm)
	}
	return aiObj
}

func RequireLLM(lm llm.LLM) map[string]any {
	return map[string]any{
		"ask": func(argsIn goja.FunctionCall, vm *goja.Runtime) goja.Value {
			args := MakeArgs(&argsIn, vm).Check(1)
			conf, cb := getAskArgs(args.This, vm, args.Arguments)
			result, usage, err := lm.Ask(makeChatMessages(args.Arguments), conf, cb)
			return makeChatResult(vm, result, &usage, err)
		},
		"fastAsk": func(argsIn goja.FunctionCall, vm *goja.Runtime) goja.Value {
			args := MakeArgs(&argsIn, vm).Check(1)
			_, cb := getAskArgs(args.This, vm, args.Arguments)
			result, usage, err := lm.FastAsk(makeChatMessages(args.Arguments), cb)
			return makeChatResult(vm, result, &usage, err)
		},
		"longAsk": func(argsIn goja.FunctionCall, vm *goja.Runtime) goja.Value {
			args := MakeArgs(&argsIn, vm).Check(1)
			_, cb := getAskArgs(args.This, vm, args.Arguments)
			result, usage, err := lm.LongAsk(makeChatMessages(args.Arguments), cb)
			return makeChatResult(vm, result, &usage, err)
		},
		"batterAsk": func(argsIn goja.FunctionCall, vm *goja.Runtime) goja.Value {
			args := MakeArgs(&argsIn, vm).Check(1)
			_, cb := getAskArgs(args.This, vm, args.Arguments)
			result, usage, err := lm.BatterAsk(makeChatMessages(args.Arguments), cb)
			return makeChatResult(vm, result, &usage, err)
		},
		"bestAsk": func(argsIn goja.FunctionCall, vm *goja.Runtime) goja.Value {
			args := MakeArgs(&argsIn, vm).Check(1)
			_, cb := getAskArgs(args.This, vm, args.Arguments)
			result, usage, err := lm.BestAsk(makeChatMessages(args.Arguments), cb)
			return makeChatResult(vm, result, &usage, err)
		},

		"multiAsk": func(argsIn goja.FunctionCall, vm *goja.Runtime) goja.Value {
			args := MakeArgs(&argsIn, vm).Check(1)
			_, cb := getAskArgs(args.This, vm, args.Arguments)
			result, usage, err := lm.MultiAsk(makeChatMessages(args.Arguments), cb)
			return makeChatResult(vm, result, &usage, err)
		},
		"bestMultiAsk": func(argsIn goja.FunctionCall, vm *goja.Runtime) goja.Value {
			args := MakeArgs(&argsIn, vm).Check(1)
			_, cb := getAskArgs(args.This, vm, args.Arguments)
			result, usage, err := lm.BestMultiAsk(makeChatMessages(args.Arguments), cb)
			return makeChatResult(vm, result, &usage, err)
		},

		"codeInterpreterAsk": func(argsIn goja.FunctionCall, vm *goja.Runtime) goja.Value {
			args := MakeArgs(&argsIn, vm).Check(1)
			_, cb := getAskArgs(args.This, vm, args.Arguments)
			result, usage, err := lm.CodeInterpreterAsk(makeChatMessages(args.Arguments), cb)
			return makeChatResult(vm, result, &usage, err)
		},
		"webSearchAsk": func(argsIn goja.FunctionCall, vm *goja.Runtime) goja.Value {
			args := MakeArgs(&argsIn, vm).Check(1)
			_, cb := getAskArgs(args.This, vm, args.Arguments)
			result, usage, err := lm.WebSearchAsk(makeChatMessages(args.Arguments), cb)
			return makeChatResult(vm, result, &usage, err)
		},

		"makeImage": func(argsIn goja.FunctionCall, vm *goja.Runtime) goja.Value {
			args := MakeArgs(&argsIn, vm).Check(1)
			prompt, conf := getGCArgs(args.Arguments)
			results, usage, err := lm.MakeImage(prompt, conf)
			return makeGCResult(vm, results, nil, &usage, err)
		},
		"fastMakeImage": func(argsIn goja.FunctionCall, vm *goja.Runtime) goja.Value {
			args := MakeArgs(&argsIn, vm).Check(1)
			prompt, conf := getGCArgs(args.Arguments)
			results, usage, err := lm.FastMakeImage(prompt, conf)
			return makeGCResult(vm, results, nil, &usage, err)
		},
		"bestMakeImage": func(argsIn goja.FunctionCall, vm *goja.Runtime) goja.Value {
			args := MakeArgs(&argsIn, vm).Check(1)
			prompt, conf := getGCArgs(args.Arguments)
			results, usage, err := lm.BestMakeImage(prompt, conf)
			return makeGCResult(vm, results, nil, &usage, err)
		},

		"makeVideo": func(argsIn goja.FunctionCall, vm *goja.Runtime) goja.Value {
			args := MakeArgs(&argsIn, vm).Check(1)
			prompt, conf := getGCArgs(args.Arguments)
			results, previews, usage, err := lm.MakeVideo(prompt, conf)
			return makeGCResult(vm, results, previews, &usage, err)
		},
		"fastMakeVideo": func(argsIn goja.FunctionCall, vm *goja.Runtime) goja.Value {
			args := MakeArgs(&argsIn, vm).Check(1)
			prompt, conf := getGCArgs(args.Arguments)
			results, previews, usage, err := lm.FastMakeVideo(prompt, conf)
			return makeGCResult(vm, results, previews, &usage, err)
		},
		"bestMakeVideo": func(argsIn goja.FunctionCall, vm *goja.Runtime) goja.Value {
			args := MakeArgs(&argsIn, vm).Check(1)
			prompt, conf := getGCArgs(args.Arguments)
			results, previews, usage, err := lm.BestMakeVideo(prompt, conf)
			return makeGCResult(vm, results, previews, &usage, err)
		},

		"embedding": func(argsIn goja.FunctionCall, vm *goja.Runtime) goja.Value {
			args := MakeArgs(&argsIn, vm).Check(2)
			results, usage, err := lm.Embedding(args.Str(0), args.Str(1))
			return makeEmbeddingResult(vm, results, &usage, err)
		},
		"fastEmbedding": func(argsIn goja.FunctionCall, vm *goja.Runtime) goja.Value {
			args := MakeArgs(&argsIn, vm).Check(1)
			results, usage, err := lm.FastEmbedding(args.Str(0))
			return makeEmbeddingResult(vm, results, &usage, err)
		},
		"bestEmbedding": func(argsIn goja.FunctionCall, vm *goja.Runtime) goja.Value {
			args := MakeArgs(&argsIn, vm).Check(1)
			results, usage, err := lm.BestEmbedding(args.Str(0))
			return makeEmbeddingResult(vm, results, &usage, err)
		},

		"support": lm.Support(),
	}
}

func getErrorStr(err error) string {
	if err != nil {
		return err.Error()
	}
	return ""
}

func makeChatResult(vm *goja.Runtime, result string, usage *llm.Usage, err error) goja.Value {
	if err != nil {
		panic(vm.NewGoError(err))
	}
	return vm.ToValue(map[string]any{
		"result":       result,
		"askTokens":    usage.AskTokens,
		"answerTokens": usage.AnswerTokens,
		"totalTokens":  usage.TotalTokens,
		"usedTime":     usage.UsedTime,
	})
}

func makeEmbeddingResult(vm *goja.Runtime, result []byte, usage *llm.Usage, err error) goja.Value {
	if err != nil {
		panic(vm.NewGoError(err))
	}
	return vm.ToValue(map[string]any{
		"result":       result,
		"askTokens":    usage.AskTokens,
		"answerTokens": usage.AnswerTokens,
		"totalTokens":  usage.TotalTokens,
		"usedTime":     usage.UsedTime,
	})
}

func makeGCResult(vm *goja.Runtime, results []string, previews []string, usage *llm.Usage, err error) goja.Value {
	if err != nil {
		panic(vm.NewGoError(err))
	}
	result := ""
	preview := ""
	if len(results) > 0 {
		result = results[0]
	} else {
		results = make([]string, 0)
	}
	if len(previews) > 0 {
		preview = previews[0]
	} else {
		previews = make([]string, 0)
	}
	return vm.ToValue(map[string]any{
		"result":   result,
		"preview":  preview,
		"results":  results,
		"previews": previews,
		"usedTime": usage.UsedTime,
	})
}

func getGCArgs(args []goja.Value) (string, llm.GCConfig) {
	prompt := ""
	var config llm.GCConfig
	if len(args) > 0 {
		prompt = u.String(args[0].Export())
		if len(args) > 1 {
			u.Convert(args[1].Export(), &config)
		}
	}
	return prompt, config
}

func getAskArgs(thisArg goja.Value, vm *goja.Runtime, args []goja.Value) (llm.ChatConfig, func(string)) {
	var chatConfig llm.ChatConfig
	var callback func(answer string)
	if len(args) > 0 {
		for i := 1; i < len(args); i++ {
			if cb, ok := goja.AssertFunction(args[i]); ok {
				callback = func(answer string) {
					_, _ = cb(thisArg, vm.ToValue(answer))
				}
			} else if args[i].ExportType() != nil {
				switch args[i].ExportType().Kind() {
				case reflect.Map, reflect.Struct:
					u.Convert(args[i].Export(), &chatConfig)
				default:
					chatConfig.Model = u.String(args[i].Export())
				}
			}
		}
	}
	return chatConfig, callback
}

func makeChatMessages(args []goja.Value) []llm.ChatMessage {
	out := make([]llm.ChatMessage, 0)
	if len(args) > 0 {
		v := args[0].Export()
		vv := reflect.ValueOf(v)
		t := args[0].ExportType()
		if t != nil {
			lastRoleIsUser := false
			switch t.Kind() {
			// 数组，根据成员类型处理
			//   字符串：
			//     含有媒体：单条多模态消息
			//     无媒体：多条文本消息
			//   数组：多条消息（第一个成员不是 role 则自动生成）
			//   对象：多条消息（无 role 则自动生成）（支持 content 或 contents）
			//   结构：转换为 llm.ChatMessage
			// 对象：单条消息（支持 content 或 contents）
			// 结构：转换为 llm.ChatMessage
			// 字符串：单条文本消息
			case reflect.Slice:
				hasSub := false
				hasMulti := false
				for i := 0; i < vv.Len(); i++ {
					vv2 := u.FinalValue(vv.Index(i))
					if vv2.Kind() == reflect.Slice || vv2.Kind() == reflect.Map || vv2.Kind() == reflect.Struct {
						hasSub = true
						break
					}
					if vv2.Kind() == reflect.String {
						str := vv2.String()
						if strings.HasPrefix(str, "data:") || strings.HasPrefix(str, "https://") || strings.HasPrefix(str, "http://") {
							hasMulti = true
						}
					}
				}

				if hasSub || !hasMulti {
					// 有子对象或纯文本数组
					var defaultRole string
					for i := 0; i < vv.Len(); i++ {
						lastRoleIsUser = !lastRoleIsUser
						if lastRoleIsUser {
							defaultRole = llm.RoleUser
						} else {
							defaultRole = llm.RoleAssistant
						}
						vv2 := u.FinalValue(vv.Index(i))
						switch vv2.Kind() {
						case reflect.Slice:
							out = append(out, makeChatMessageFromSlice(vv2, defaultRole))
						case reflect.Map:
							out = append(out, makeChatMessageFromSlice(vv2, defaultRole))
						case reflect.Struct:
							item := llm.ChatMessage{}
							u.Convert(vv2.Interface(), &item)
							out = append(out, item)
						default:
							out = append(out, llm.ChatMessage{Role: llm.RoleUser, Contents: []llm.ChatMessageContent{makeChatMessageContent(u.String(vv2.Interface()))}})
						}
						lastRoleIsUser = out[len(out)-1].Role != llm.RoleUser
					}
				} else {
					// 单条多模态消息
					out = append(out, makeChatMessageFromSlice(vv, llm.RoleUser))
				}
			case reflect.Map:
				out = append(out, makeChatMessageFromMap(vv, llm.RoleUser))
			case reflect.Struct:
				item := llm.ChatMessage{}
				u.Convert(v, &item)
				out = append(out, item)
			default:
				out = append(out, llm.ChatMessage{Role: llm.RoleUser, Contents: []llm.ChatMessageContent{makeChatMessageContent(u.String(v))}})
			}
		}
	}
	return out
}

func makeChatMessageFromSlice(vv reflect.Value, defaultRole string) llm.ChatMessage {
	role := u.String(vv.Index(0).Interface())
	j := 0
	if role == llm.RoleUser || role == llm.RoleAssistant || role == llm.RoleSystem || role == llm.RoleTool {
		j = 1
	} else {
		role = defaultRole
	}
	contents := make([]llm.ChatMessageContent, 0)
	for ; j < vv.Len(); j++ {
		contents = append(contents, makeChatMessageContent(u.String(vv.Index(j).Interface())))
	}
	return llm.ChatMessage{Role: role, Contents: contents}
}

func makeChatMessageFromMap(vv reflect.Value, defaultRole string) llm.ChatMessage {
	role := u.String(vv.MapIndex(reflect.ValueOf("role")).Interface())
	if role == "" {
		role = defaultRole
	}
	contents := make([]llm.ChatMessageContent, 0)
	content := u.String(vv.MapIndex(reflect.ValueOf("content")).Interface())
	if content != "" {
		contents = append(contents, makeChatMessageContent(content))
	} else {
		contentsV := vv.MapIndex(reflect.ValueOf("contents"))
		if contentsV.IsValid() && contentsV.Kind() == reflect.Slice {
			for i := 0; i < contentsV.Len(); i++ {
				contents = append(contents, makeChatMessageContent(u.String(contentsV.Index(i).Interface())))
			}
		}
	}
	return llm.ChatMessage{Role: role, Contents: contents}
}

func makeChatMessageContent(contnet string) llm.ChatMessageContent {
	if strings.HasPrefix(contnet, "data:image/") || ((strings.HasPrefix(contnet, "https://") || strings.HasPrefix(contnet, "http://")) && (strings.HasSuffix(contnet, ".png") || strings.HasSuffix(contnet, ".jpg") || strings.HasSuffix(contnet, ".jpeg") || strings.HasSuffix(contnet, ".gif") || strings.HasSuffix(contnet, ".svg"))) {
		return llm.ChatMessageContent{Type: llm.TypeImage, Content: contnet}
	} else if strings.HasPrefix(contnet, "data:video/") || ((strings.HasPrefix(contnet, "https://") || strings.HasPrefix(contnet, "http://")) && (strings.HasSuffix(contnet, ".mp4") || strings.HasSuffix(contnet, ".mov") || strings.HasSuffix(contnet, ".m4v") || strings.HasSuffix(contnet, ".avi") || strings.HasSuffix(contnet, ".wmv"))) {
		return llm.ChatMessageContent{Type: llm.TypeVideo, Content: contnet}
	}
	return llm.ChatMessageContent{Type: llm.TypeText, Content: contnet}
}