package ernie

import (
	"gitee.com/kzangv/gsf-ai-agent/llms"
)

/**
const
*/

const (
	_BaiduBceURL                    = "https://aip.baidubce.com"
	_DefaultEmptyMessagesLimit uint = 300
	TokenPrefix                     = "baidu:token:"

	defaultTemperature  = 0.95
	defaultTopP         = 0.8
	defaultPenaltyScore = 1
	defaultMaxTokens    = 500
)

type ModelType int

const (
	ModelTypeErnieBot    ModelType = 1 // 对应模型 ERNIE-Bot
	ModelTypeErnieBot8K  ModelType = 2 // 对应模型 ERNIE-Bot-8K
	ModelTypeErnieBotPro ModelType = 3 // 对应模型 ERNIE-Bot 4.0
)

const (
	_MessageRoleUser      = "user"
	_MessageRoleAssistant = "assistant"
	_MessageRoleFunction  = "function"
)

const (
	_FinishNormal        = "normal"         // 输出内容完全由大模型生成，未触发截断、替换
	_FinishStop          = "stop"           // 输出结果命中入参 stop 中指定的字段后被截断
	_FinishLength        = "length"         // 达到了最大的 token 数，根据 EB 返回结果 is_truncated 来截断
	_FinishContentFilter = "content_filter" // 输出内容被截断、兜底、替换为 ** 等
	_FinishFunctionCall  = "function_call"  // 调用了 function call功能
)

var (
	_ModelReqUrlMap = map[ModelType]string{
		ModelTypeErnieBot:    "/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/completions",
		ModelTypeErnieBot8K:  "/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/ernie_bot_8k",
		ModelTypeErnieBotPro: "/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/completions_pro",
	}

	_ModelRespFinishMap = map[string]llms.OutputFinishType{
		_FinishNormal:        llms.OutputFinishNormal,
		_FinishStop:          llms.OutputFinishStop,
		_FinishLength:        llms.OutputFinishLength,
		_FinishContentFilter: llms.OutputFinishContentFilter,
		_FinishFunctionCall:  llms.OutputFinishToolCalls,
	}
)

/**
Common
*/

type ErnieFunctionCall struct {
	Name      string `json:"name"`
	Arguments string `json:"arguments"`
	Thoughts  string `json:"thoughts"`
}

type ErnieFunctionExample struct {
	Role         string             `json:"role"`
	Content      *string            `json:"content"`
	Name         string             `json:"name,omitempty"`
	FunctionCall *ErnieFunctionCall `json:"function_call,omitempty"`
}

type ErnieFunction struct {
	Name        string                  `json:"name"`
	Description string                  `json:"description"`
	Parameters  any                     `json:"parameters"`
	Responses   any                     `json:"responses"`
	Examples    []*ErnieFunctionExample `json:"examples,omitempty"`
}

/**
Request
*/

type ReqChatCompMsg struct {
	Role         string             `json:"role"`
	Content      *string            `json:"content"`
	Name         string             `json:"name,omitempty"`
	FunctionCall *ErnieFunctionCall `json:"function_call,omitempty"`
}

type Request struct {
	Messages        []*ReqChatCompMsg `json:"messages"`
	Temperature     float32           `json:"temperature,omitempty"`
	TopP            float32           `json:"top_p,omitempty"`
	PresencePenalty float32           `json:"presence_penalty,omitempty"`
	Stream          bool              `json:"stream"`
	UserId          string            `json:"user_id,omitempty"`
	PenaltyScore    float32           `json:"penalty_score,omitempty"`
	Functions       []*ErnieFunction  `json:"functions,omitempty"`
	System          string            `json:"system,omitempty"`
	Stop            []string          `json:"stop,omitempty"`
	DisableSearch   bool              `json:"disable_search,omitempty"`
	EnableCitation  bool              `json:"enable_citation,omitempty"`
}

/**
Response
*/

type RespError struct {
	ErrorCode int    `json:"error_code"`
	ErrorMsg  string `json:"error_msg"`
}

type RespPluginUsage struct {
	Name           string `json:"name"`
	ParseTokens    int    `json:"parse_tokens"`
	AbstractTokens int    `json:"abstract_tokens"`
	SearchTokens   int    `json:"search_tokens"`
	TotalTokens    int    `json:"total_tokens"`
}

type RespUsage struct {
	PromptTokens     int                `json:"prompt_tokens"`
	CompletionTokens int                `json:"completion_tokens"`
	TotalTokens      int                `json:"total_tokens"`
	Plugins          []*RespPluginUsage `json:"plugins,omitempty"`
}

type RespSearchResult struct {
	Index        int    `json:"index"`
	Url          string `json:"url"`
	Title        string `json:"title"`
	DataSourceId string `json:"data_source_id"`
}

type RespSearchInfo struct {
	IsBeset       bool                `json:"is_beset"`
	ReWrite       string              `json:"re_write"`
	SearchResults []*RespSearchResult `json:"search_results,omitempty"`
}

type Response struct {
	RespError

	ID               string             `json:"id"`
	Object           string             `json:"object"`
	Created          int                `json:"created"`
	SentenceId       int                `json:"sentence_id"`
	BanRound         int                `json:"ban_round"`
	IsEnd            bool               `json:"is_end"`
	IsTruncated      bool               `json:"is_truncated"`
	NeedClearHistory bool               `json:"need_clear_history"`
	FinishReason     string             `json:"finish_reason"`
	Result           string             `json:"result"`
	Usage            RespUsage          `json:"usage"`
	FunctionCall     *ErnieFunctionCall `json:"function_call,omitempty"`
	SearchInfo       *RespSearchInfo    `json:"search_info,omitempty"`
}

type ResponseAccessToken struct {
	RefreshToken     string `json:"refresh_token"`
	ExpiresIn        int    `json:"expires_in"`
	SessionKey       string `json:"session_key"`
	AccessToken      string `json:"access_token"`
	Scope            string `json:"scope"`
	SessionSecret    string `json:"session_secret"`
	Error            string `json:"error"`
	ErrorDescription string `json:"error_description"`
}