package agent

import "sync"

const (
	TypeText            = "text"
	TypeImage           = "image"
	TypeVideo           = "video"
	RoleSystem          = "system"
	RoleUser            = "user"
	RoleAssistant       = "assistant"
	RoleTool            = "tool"
	ToolCodeInterpreter = "codeInterpreter"
	ToolWebSearch       = "webSearch"
)

type Support struct {
	Ask                    bool
	AskWithImage           bool
	AskWithVideo           bool
	AskWithCodeInterpreter bool
	AskWithWebSearch       bool
	MakeImage              bool
	MakeVideo              bool
	Models                 []string
}

type Agent interface {
	Support() Support
	Ask(messages []ChatMessage, config *ChatModelConfig, callback func(answer string)) (string, TokenUsage, error)
	FastAsk(messages []ChatMessage, callback func(answer string)) (string, TokenUsage, error)
	LongAsk(messages []ChatMessage, callback func(answer string)) (string, TokenUsage, error)
	BatterAsk(messages []ChatMessage, callback func(answer string)) (string, TokenUsage, error)
	BestAsk(messages []ChatMessage, callback func(answer string)) (string, TokenUsage, error)
	MultiAsk(messages []ChatMessage, callback func(answer string)) (string, TokenUsage, error)
	BestMultiAsk(messages []ChatMessage, callback func(answer string)) (string, TokenUsage, error)
	CodeInterpreterAsk(messages []ChatMessage, callback func(answer string)) (string, TokenUsage, error)
	WebSearchAsk(messages []ChatMessage, callback func(answer string)) (string, TokenUsage, error)
	MakeImage(model, prompt, size, refImage string) ([]string, error)
	FastMakeImage(prompt, size, refImage string) ([]string, error)
	BestMakeImage(prompt, size, refImage string) ([]string, error)
	MakeVideo(model, prompt, size, refImage string) ([]string, []string, error)
	FastMakeVideo(prompt, size, refImage string) ([]string, []string, error)
	BestMakeVideo(prompt, size, refImage string) ([]string, []string, error)
}

type APIConfig struct {
	Endpoint               string
	ApiKey                 string
	DefaultChatModelConfig ChatModelConfig
}

type TokenUsage struct {
	AskTokens    int64
	AnswerTokens int64
	TotalTokens  int64
}

var agentMakers = map[string]func(APIConfig) Agent{}
var agentMakersLock = sync.RWMutex{}

var agents = map[string]Agent{}
var agentsLock = sync.RWMutex{}

func RegisterAgentMaker(agentId string, maker func(APIConfig) Agent) {
	agentMakersLock.Lock()
	agentMakers[agentId] = maker
	agentMakersLock.Unlock()
}

func CreateAgent(name, agentId string, config APIConfig) Agent {
	agentMakersLock.RLock()
	maker := agentMakers[agentId]
	agentMakersLock.RUnlock()

	if maker != nil {
		agent := maker(config)
		agentsLock.Lock()
		agents[name] = agent
		agentsLock.Unlock()
		return agent
	}
	return nil
}

func GetAgent(name string) Agent {
	agentsLock.RLock()
	agent := agents[name]
	agentsLock.RUnlock()
	return agent
}