package llm

import (
	"bytes"
	_ "embed"
	"text/template"

	"apigo.cc/ai/llm/llm"
	"apigo.cc/gojs"
	"apigo.cc/gojs/goja"
	"github.com/ssgo/config"
	"github.com/ssgo/u"
)

//go:embed llm.ts
var llmTS string

//go:embed llm.md
var llmMD string

var confAes = u.NewAes([]byte("?GQ$0K0GgLdO=f+~L68PLm$uhKr4'=tV"), []byte("VFs7@sK61cj^f?HZ"))
var keysIsSet = false

func SetSSKey(key, iv []byte) {
	if !keysIsSet {
		confAes = u.NewAes(key, iv)
		keysIsSet = true
	}
}

var jsObj gojs.Map
var llmList = map[string]*map[string]*struct {
	Endpoint   string
	ApiKey     string
	ChatConfig llm.ChatConfig
	GCConfig   llm.GCConfig
	Debug      bool
}{}

func init() {
	llmNames := make([]string, 0)
	_ = config.LoadConfig("llm", &llmList)
	for llmName, llmConfigs := range llmList {
		for confName, _ := range *llmConfigs {
			if confName == "default" {
				confName = llmName
			}
			llmNames = append(llmNames, confName)
		}
	}

	var tpl *template.Template
	var err error
	llmTSCode := ""
	if tpl, err = template.New("").Parse(llmTS); err == nil {
		buf := bytes.NewBuffer(make([]byte, 0))
		if err = tpl.Execute(buf, llmNames); err == nil {
			llmTSCode = string(buf.Bytes())
		}
	}

	gojs.Register("apigo.cc/ai/llm", gojs.Module{
		ObjectMaker: func(vm *goja.Runtime) gojs.Map {
			if jsObj == nil {
				makeJsObj()
			}
			return jsObj
		},
		TsCode:  llmTSCode,
		Desc:    "llm plugin for gojs(http://apigo.cc/gojs)",
		Example: llmMD,
	})
}

func makeJsObj() {
	jsObj = make(gojs.Map)
	for llmName, llmConfigs := range llmList {
		for confName, llmConf := range *llmConfigs {
			llmConf.ApiKey = confAes.DecryptUrlBase64ToString(llmConf.ApiKey)
			if confName == "default" {
				confName = llmName
			}
			llmObj := llm.Create(confName, llmName, llm.Config{
				Endpoint:   llmConf.Endpoint,
				ApiKey:     llmConf.ApiKey,
				ChatConfig: llmConf.ChatConfig,
				GCConfig:   llmConf.GCConfig,
				Debug:      llmConf.Debug,
			})
			if llmObj != nil {
				jsObj[confName] = MakeLLM(llmObj)
			}
		}
	}
	jsObj["similarity"] = func(argsIn goja.FunctionCall, vm *goja.Runtime) goja.Value {
		args := gojs.MakeArgs(&argsIn, vm).Check(2)
		return vm.ToValue(llm.Similarity(args.Bytes(0), args.Bytes(1)))
	}
}