package zhipu

import (
	"apigo.cc/ai/ai/interface/llm"
	"apigo.cc/ai/ai/llm/zhipu/zhipu"
)

type LLM struct {
	config llm.Config
}

var NameMap = map[string]string{
	llm.TypeText:      zhipu.MultiContentTypeText,
	llm.TypeImage:     zhipu.MultiContentTypeImageURL,
	llm.TypeVideo:     zhipu.MultiContentTypeVideoURL,
	llm.RoleSystem:    zhipu.RoleSystem,
	llm.RoleUser:      zhipu.RoleUser,
	llm.RoleAssistant: zhipu.RoleAssistant,
	llm.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 (lm *LLM) Support() llm.Support {
	return llm.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() {
	llm.Register("zhipu", func(config llm.Config) llm.LLM {
		return &LLM{config: config}
	})
}