package js

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

type ChatResult struct {
	llm.TokenUsage
	Result string
	Error  string
}

type AIGCResult struct {
	Result   string
	Preview  string
	Results  []string
	Previews []string
	Error    string
}

func RequireAI(lm llm.LLM) map[string]any {
	return map[string]any{
		"ask": func(args goja.FunctionCall, vm *goja.Runtime) goja.Value {
			conf, cb := getAskArgs(args.This, vm, args.Arguments)
			result, usage, err := lm.Ask(makeChatMessages(args.Arguments), conf, cb)
			return vm.ToValue(map[string]any{"tokenUsage": toMap(usage), "result": result, "error": getErrorStr(err)})
		},
		"fastAsk": func(args goja.FunctionCall, vm *goja.Runtime) goja.Value {
			_, cb := getAskArgs(args.This, vm, args.Arguments)
			result, usage, err := lm.FastAsk(makeChatMessages(args.Arguments), cb)
			return vm.ToValue(map[string]any{"tokenUsage": toMap(usage), "result": result, "error": getErrorStr(err)})
		},
		"longAsk": func(args goja.FunctionCall, vm *goja.Runtime) goja.Value {
			_, cb := getAskArgs(args.This, vm, args.Arguments)
			result, usage, err := lm.LongAsk(makeChatMessages(args.Arguments), cb)
			return vm.ToValue(map[string]any{"tokenUsage": toMap(usage), "result": result, "error": getErrorStr(err)})
		},
		"batterAsk": func(args goja.FunctionCall, vm *goja.Runtime) goja.Value {
			_, cb := getAskArgs(args.This, vm, args.Arguments)
			result, usage, err := lm.BatterAsk(makeChatMessages(args.Arguments), cb)
			return vm.ToValue(map[string]any{"tokenUsage": toMap(usage), "result": result, "error": getErrorStr(err)})
		},
		"bestAsk": func(args goja.FunctionCall, vm *goja.Runtime) goja.Value {
			_, cb := getAskArgs(args.This, vm, args.Arguments)
			result, usage, err := lm.BestAsk(makeChatMessages(args.Arguments), cb)
			return vm.ToValue(map[string]any{"tokenUsage": toMap(usage), "result": result, "error": getErrorStr(err)})
		},

		"multiAsk": func(args goja.FunctionCall, vm *goja.Runtime) goja.Value {
			_, cb := getAskArgs(args.This, vm, args.Arguments)
			result, usage, err := lm.MultiAsk(makeChatMessages(args.Arguments), cb)
			return vm.ToValue(map[string]any{"tokenUsage": toMap(usage), "result": result, "error": getErrorStr(err)})
		},
		"bestMultiAsk": func(args goja.FunctionCall, vm *goja.Runtime) goja.Value {
			_, cb := getAskArgs(args.This, vm, args.Arguments)
			result, usage, err := lm.BestMultiAsk(makeChatMessages(args.Arguments), cb)
			return vm.ToValue(map[string]any{"tokenUsage": toMap(usage), "result": result, "error": getErrorStr(err)})
		},

		"codeInterpreterAsk": func(args goja.FunctionCall, vm *goja.Runtime) goja.Value {
			_, cb := getAskArgs(args.This, vm, args.Arguments)
			result, usage, err := lm.CodeInterpreterAsk(makeChatMessages(args.Arguments), cb)
			return vm.ToValue(map[string]any{"tokenUsage": toMap(usage), "result": result, "error": getErrorStr(err)})
		},
		"webSearchAsk": func(args goja.FunctionCall, vm *goja.Runtime) goja.Value {
			_, cb := getAskArgs(args.This, vm, args.Arguments)
			result, usage, err := lm.WebSearchAsk(makeChatMessages(args.Arguments), cb)
			return vm.ToValue(map[string]any{"tokenUsage": toMap(usage), "result": result, "error": getErrorStr(err)})
		},

		"makeImage": func(args goja.FunctionCall, vm *goja.Runtime) goja.Value {
			prompt, conf := getAIGCArgs(args.Arguments)
			results, err := lm.MakeImage(prompt, conf)
			return makeAIGCResult(vm, results, nil, err)
		},
		"fastMakeImage": func(args goja.FunctionCall, vm *goja.Runtime) goja.Value {
			prompt, conf := getAIGCArgs(args.Arguments)
			results, err := lm.FastMakeImage(prompt, conf)
			return makeAIGCResult(vm, results, nil, err)
		},
		"bestMakeImage": func(args goja.FunctionCall, vm *goja.Runtime) goja.Value {
			prompt, conf := getAIGCArgs(args.Arguments)
			results, err := lm.BestMakeImage(prompt, conf)
			return makeAIGCResult(vm, results, nil, err)
		},

		"makeVideo": func(args goja.FunctionCall, vm *goja.Runtime) goja.Value {
			prompt, conf := getAIGCArgs(args.Arguments)
			results, previews, err := lm.MakeVideo(prompt, conf)
			return makeAIGCResult(vm, results, previews, err)
		},
		"fastMakeVideo": func(args goja.FunctionCall, vm *goja.Runtime) goja.Value {
			prompt, conf := getAIGCArgs(args.Arguments)
			results, previews, err := lm.FastMakeVideo(prompt, conf)
			return makeAIGCResult(vm, results, previews, err)
		},
		"bestMakeVideo": func(args goja.FunctionCall, vm *goja.Runtime) goja.Value {
			prompt, conf := getAIGCArgs(args.Arguments)
			results, previews, err := lm.BestMakeVideo(prompt, conf)
			return makeAIGCResult(vm, results, previews, err)
		},

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

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

func makeAIGCResult(vm *goja.Runtime, results []string, previews []string, err error) goja.Value {
	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,
		"error":    getErrorStr(err),
	})
}

func getAIGCArgs(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 {
				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()
		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}
}