package ai

import (
	"apigo.cc/ai/agent"
	"github.com/ssgo/config"
	"github.com/ssgo/u"
	"sync"
)

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
	}
}

const (
	TypeText            = agent.TypeText
	TypeImage           = agent.TypeImage
	TypeVideo           = agent.TypeVideo
	RoleSystem          = agent.RoleSystem
	RoleUser            = agent.RoleUser
	RoleAssistant       = agent.RoleAssistant
	RoleTool            = agent.RoleTool
	ToolCodeInterpreter = agent.ToolCodeInterpreter
	ToolWebSearch       = agent.ToolWebSearch
)

type APIConfig struct {
	Endpoint               string
	ApiKey                 string
	DefaultChatModelConfig ChatModelConfig
}

type AgentConfig struct {
	ApiKey     string
	Endpoint   string
	Agent      string
	ChatConfig ChatModelConfig
}

var agentConfigs map[string]*AgentConfig
var agentConfigsLock = sync.RWMutex{}

func GetAgent(name string) agent.Agent {
	ag := agent.GetAgent(name)
	if ag != nil {
		return ag
	}

	var agConf *AgentConfig
	if agentConfigs == nil {
		agConfs := make(map[string]*AgentConfig)
		config.LoadConfig("agent", &agConfs)
		agConf = agConfs[name]
		if agConf != nil {
			agentConfigsLock.Lock()
			agentConfigs = agConfs
			agentConfigsLock.Unlock()
		}
	} else {
		agentConfigsLock.RLock()
		agConf = agentConfigs[name]
		agentConfigsLock.RUnlock()
	}

	if agConf == nil {
		return nil
	}

	if agConf.Agent == "" {
		agConf.Agent = name
	}

	return agent.CreateAgent(name, agConf.Agent, agent.APIConfig{
		Endpoint: agConf.Endpoint,
		ApiKey:   confAes.DecryptUrlBase64ToString(agConf.ApiKey),
		DefaultChatModelConfig: agent.ChatModelConfig{
			Model:       agConf.ChatConfig.Model,
			MaxTokens:   agConf.ChatConfig.MaxTokens,
			Temperature: agConf.ChatConfig.Temperature,
			TopP:        agConf.ChatConfig.TopP,
			Tools:       agConf.ChatConfig.Tools,
		},
	})
}