package openai

import (
	"apigo.cc/ai/ai/llm"
	"github.com/sashabaranov/go-openai"
)

type LLM struct {
	config llm.Config
}

var TypeMap = map[string]openai.ChatMessagePartType{
	llm.TypeText:  openai.ChatMessagePartTypeText,
	llm.TypeImage: openai.ChatMessagePartTypeImageURL,
	//llm.TypeVideo: "video_url",
}
var RoleMap = map[string]string{
	llm.RoleSystem:    openai.ChatMessageRoleSystem,
	llm.RoleUser:      openai.ChatMessageRoleUser,
	llm.RoleAssistant: openai.ChatMessageRoleAssistant,
	llm.RoleTool:      openai.ChatMessageRoleTool,
}

const (
	ModelGPT_4_32k_0613         = "gpt-4-32k-0613"
	ModelGPT_4_32k_0314         = "gpt-4-32k-0314"
	ModelGPT_4_32k              = "gpt-4-32k"
	ModelGPT_4_0613             = "gpt-4-0613"
	ModelGPT_4_0314             = "gpt-4-0314"
	ModelGPT_4o                 = "gpt-4o"
	ModelGPT_4o_2024_05_13      = "gpt-4o-2024-05-13"
	ModelGPT_4o_2024_08_06      = "gpt-4o-2024-08-06"
	ModelGPT_4o_mini            = "gpt-4o-mini"
	ModelGPT_4o_mini_2024_07_18 = "gpt-4o-mini-2024-07-18"
	ModelGPT_4_turbo            = "gpt-4-turbo"
	ModelGPT_4_turbo_2024_04_09 = "gpt-4-turbo-2024-04-09"
	ModelGPT_4_0125_preview     = "gpt-4-0125-preview"
	ModelGPT_4_1106_preview     = "gpt-4-1106-preview"
	ModelGPT_4_turbo_preview    = "gpt-4-turbo-preview"
	ModelGPT_4_vision_preview   = "gpt-4-vision-preview"
	ModelGPT_4                  = "gpt-4"
	ModelGPT_3_5_turbo_0125     = "gpt-3.5-turbo-0125"
	ModelGPT_3_5_turbo_1106     = "gpt-3.5-turbo-1106"
	ModelGPT_3_5_turbo_0613     = "gpt-3.5-turbo-0613"
	ModelGPT_3_5_turbo_0301     = "gpt-3.5-turbo-0301"
	ModelGPT_3_5_turbo_16k      = "gpt-3.5-turbo-16k"
	ModelGPT_3_5_turbo_16k_0613 = "gpt-3.5-turbo-16k-0613"
	ModelGPT_3_5_turbo          = "gpt-3.5-turbo"
	ModelGPT_3_5_turbo_instruct = "gpt-3.5-turbo-instruct"
	ModelDavinci_002            = "davinci-002"
	ModelCurie                  = "curie"
	ModelCurie_002              = "curie-002"
	ModelAda_002                = "ada-002"
	ModelBabbage_002            = "babbage-002"
	ModelCode_davinci_002       = "code-davinci-002"
	ModelCode_cushman_001       = "code-cushman-001"
	ModelCode_davinci_001       = "code-davinci-001"
	ModelDallE2Std              = "dall-e-2"
	ModelDallE2HD               = "dall-e-2-hd"
	ModelDallE3Std              = "dall-e-3"
	ModelDallE3HD               = "dall-e-3-hd"
)

func (ag *LLM) Support() llm.Support {
	return llm.Support{
		Ask:                    true,
		AskWithImage:           true,
		AskWithVideo:           false,
		AskWithCodeInterpreter: true,
		AskWithWebSearch:       false,
		MakeImage:              true,
		MakeVideo:              false,
		Models:                 []string{ModelGPT_4_32k_0613, ModelGPT_4_32k_0314, ModelGPT_4_32k, ModelGPT_4_0613, ModelGPT_4_0314, ModelGPT_4o, ModelGPT_4o_2024_05_13, ModelGPT_4o_2024_08_06, ModelGPT_4o_mini, ModelGPT_4o_mini_2024_07_18, ModelGPT_4_turbo, ModelGPT_4_turbo_2024_04_09, ModelGPT_4_0125_preview, ModelGPT_4_1106_preview, ModelGPT_4_turbo_preview, ModelGPT_4_vision_preview, ModelGPT_4, ModelGPT_3_5_turbo_0125, ModelGPT_3_5_turbo_1106, ModelGPT_3_5_turbo_0613, ModelGPT_3_5_turbo_0301, ModelGPT_3_5_turbo_16k, ModelGPT_3_5_turbo_16k_0613, ModelGPT_3_5_turbo, ModelGPT_3_5_turbo_instruct, ModelDavinci_002, ModelCurie, ModelCurie_002, ModelAda_002, ModelBabbage_002, ModelCode_davinci_002, ModelCode_cushman_001, ModelCode_davinci_001, ModelDallE2Std, ModelDallE2HD, ModelDallE3Std, ModelDallE3HD},
	}
}

func init() {
	llm.Register("openai", func(config llm.Config) llm.LLM {
		return &LLM{config: config}
	})
}