package zhipu

import (
	"apigo.cc/ai/agent"
	"apigo.cc/ai/zhipu/zhipu"
)

type Agent struct {
	config agent.APIConfig
}

var NameMap = map[string]string{
	agent.TypeText:      zhipu.MultiContentTypeText,
	agent.TypeImage:     zhipu.MultiContentTypeImageURL,
	agent.TypeVideo:     zhipu.MultiContentTypeVideoURL,
	agent.RoleSystem:    zhipu.RoleSystem,
	agent.RoleUser:      zhipu.RoleUser,
	agent.RoleAssistant: zhipu.RoleAssistant,
	agent.RoleTool:      zhipu.RoleTool,
}

const (
	ModelGLM4Plus     = "GLM-4-Plus"
	ModelGLM40520     = "GLM-4-0520"
	ModelGLM4Long     = "GLM-4-Long"
	ModelGLM4AirX     = "GLM-4-AirX"
	ModelGLM4Air      = "GLM-4-Air"
	ModelGLM4Flash    = "GLM-4-Flash"
	ModelGLM4AllTools = "GLM-4-AllTools"
	ModelGLM4         = "GLM-4"
	ModelGLM4VPlus    = "GLM-4V-Plus"
	ModelGLM4V        = "GLM-4V"
	ModelCogVideoX    = "CogVideoX"
	ModelCogView3Plus = "CogView-3-Plus"
	ModelCogView3     = "CogView-3"
	ModelEmbedding3   = "Embedding-3"
	ModelEmbedding2   = "Embedding-2"
	ModelCharGLM3     = "CharGLM-3"
	ModelEmohaa       = "Emohaa"
	ModelCodeGeeX4    = "CodeGeeX-4"
)

func (ag *Agent) Support() agent.Support {
	return agent.Support{
		Ask:                    true,
		AskWithImage:           true,
		AskWithVideo:           true,
		AskWithCodeInterpreter: true,
		AskWithWebSearch:       true,
		MakeImage:              true,
		MakeVideo:              true,
		Models:                 []string{ModelGLM4Plus, ModelGLM40520, ModelGLM4Long, ModelGLM4AirX, ModelGLM4Air, ModelGLM4Flash, ModelGLM4AllTools, ModelGLM4, ModelGLM4VPlus, ModelGLM4V, ModelCogVideoX, ModelCogView3Plus, ModelCogView3, ModelEmbedding3, ModelEmbedding2, ModelCharGLM3, ModelEmohaa, ModelCodeGeeX4},
	}
}

func init() {
	agent.RegisterAgentMaker("zhipu", func(config agent.APIConfig) agent.Agent {
		return &Agent{config: config}
	})
}