package huoshan

import (
	"strings"

	"apigo.cc/ai/llm/llm"
	"github.com/volcengine/volc-sdk-golang/service/visual"
	"github.com/volcengine/volcengine-go-sdk/service/arkruntime"
	"github.com/volcengine/volcengine-go-sdk/service/arkruntime/model"
)

type LLM struct {
	config llm.Config
}

var NameMap = map[string]string{
	llm.TypeText:  string(model.ChatCompletionMessageContentPartTypeText),
	llm.TypeImage: string(model.ChatCompletionMessageContentPartTypeImageURL),
	//llm.TypeVideo:     string(model.ChatCompletionMessageContentPartTypeVideoURL),
	llm.RoleSystem:    model.ChatMessageRoleSystem,
	llm.RoleUser:      model.ChatMessageRoleUser,
	llm.RoleAssistant: model.ChatMessageRoleAssistant,
	llm.RoleTool:      model.ChatMessageRoleTool,
}

const (
	ModelDoubaoLite4k                 = "Doubao-lite-4k"
	ModelDoubaoLite32k                = "Doubao-lite-32k"
	ModelDoubaoLite128k               = "Doubao-lite-128k"
	ModelDoubaoPro4k                  = "Doubao-pro-4k"
	ModelDoubaoPro32k                 = "Doubao-pro-32k"
	ModelDoubaoPro128k                = "Doubao-pro-128k"
	ModelDoubaoPro256k                = "Doubao-pro-256k"
	ModelDoubaoEmbedding              = "Doubao-embedding"
	ModelDoubaoEmbeddingLarge         = "Doubao-embedding-large"
	ModelT2I2L                        = "high_aes_general_v20_L:general_v2.0_L"
	ModelT2I2S                        = "high_aes_general_v20:general_v2.0"
	ModelT2IXL                        = "t2i_xl_sft"
	ModelI2IXL                        = "i2i_xl_sft"
	ModelT2I14                        = "high_aes_general_v14"
	ModelI2I14IP                      = "high_aes_general_v14_ip_keep"
	ModelAnime13                      = "high_aes:anime_v1.3"
	ModelAnime131                     = "high_aes:anime_v1.3.1"
	ModelPhotoverseAmericanComics     = "img2img_photoverse_american_comics"    // 美漫风格
	ModelPhotoverseExecutiveIDPhoto   = "img2img_photoverse_executive_ID_photo" // 商务证件照
	ModelPhotoverse3dWeird            = "img2img_photoverse_3d_weird"           // 3d人偶
	ModelPhotoverseCyberpunk          = "img2img_photoverse_cyberpunk"          // 赛博朋克
	ModelXiezhenGubao                 = "img2img_xiezhen_gubao"                 // 古堡
	ModelXiezhenBabiNiuzai            = "img2img_xiezhen_babi_niuzai"           // 芭比牛仔
	ModelXiezhenBathrobe              = "img2img_xiezhen_bathrobe"              // 浴袍风格
	ModelXiezhenButterflyMachin       = "img2img_xiezhen_butterfly_machin"      // 蝴蝶机械
	ModelXiezhenZhichangzhengjianzhao = "img2img_xiezhen_zhichangzhengjianzhao" // 职场证件照
	ModelXiezhenChristmas             = "img2img_xiezhen_christmas"             // 圣诞
	ModelXiezhenDessert               = "img2img_xiezhen_dessert"               // 美式甜点师
	ModelXiezhenOldMoney              = "img2img_xiezhen_old_money"             // old money
	ModelXiezhenSchool                = "img2img_xiezhen_school"                // 最美校园
)

func (lm *LLM) Support() llm.Support {
	return llm.Support{
		Ask:                    true,
		AskWithImage:           true,
		AskWithVideo:           false,
		AskWithCodeInterpreter: false,
		AskWithWebSearch:       false,
		MakeImage:              true,
		MakeVideo:              false,
		Models:                 []string{ModelDoubaoLite4k, ModelDoubaoLite32k, ModelDoubaoLite128k, ModelDoubaoPro4k, ModelDoubaoPro32k, ModelDoubaoPro128k, ModelDoubaoPro256k, ModelDoubaoEmbedding, ModelDoubaoEmbeddingLarge, ModelT2I2L, ModelT2I2S, ModelT2IXL, ModelI2IXL, ModelT2I14, ModelI2I14IP, ModelAnime13, ModelAnime131, ModelPhotoverseAmericanComics, ModelPhotoverseExecutiveIDPhoto, ModelPhotoverse3dWeird, ModelPhotoverseCyberpunk, ModelXiezhenGubao, ModelXiezhenBabiNiuzai, ModelXiezhenBathrobe, ModelXiezhenButterflyMachin, ModelXiezhenZhichangzhengjianzhao, ModelXiezhenChristmas, ModelXiezhenDessert, ModelXiezhenOldMoney, ModelXiezhenSchool},
	}
}

func (lm *LLM) getChatClient() *arkruntime.Client {
	opt := make([]arkruntime.ConfigOption, 0)
	if lm.config.Endpoint != "" {
		opt = append(opt, arkruntime.WithBaseUrl(lm.config.Endpoint))
	}
	return arkruntime.NewClientWithAkSk(strings.SplitN(lm.config.ApiKey, ",", 2)[0], opt...)
}

func (lm *LLM) getGCClient() *visual.Visual {
	keys := strings.SplitN(lm.config.ApiKey, ",", 2)
	if len(keys) == 1 {
		keys = append(keys, "")
	}
	vis := visual.NewInstance()
	vis.Client.SetAccessKey(keys[0])
	vis.Client.SetSecretKey(keys[1])
	return vis
}

// 因为火山平台的配置过于繁琐（每个模型都要创建单独的endpoint，所以暂时放弃对豆包大模型的支持）
// func init() {
// 	llm.Register("huoshan", func(config llm.Config) llm.LLM {
// 		return &LLM{config: config}
// 	})
// }