package doubao

import (
	"fmt"
)

type Client struct {
	enableSearch   bool    // 是否开启搜索
	searchStrategy string  // 搜索策略（开了搜索才生效）
	temperature    float32 // 采样温度
}

func NewAIClient(enableSearch bool, searchStrategy string, temperature float32) *Client {
	return &Client{
		enableSearch:   enableSearch,
		searchStrategy: searchStrategy,
		temperature:    temperature,
	}
}

func (c *Client) RequestAI(model, key string, pre []string, prompt []string, cb func(content string, i int) error) ([]Message, int32, int32, error) {
	inToken, outToken := int32(0), int32(0)
	messages := make([]Message, 0)
	for _, s := range pre {
		messages = append(messages, Message{
			Role:    "user",
			Content: s,
		})
	}
	for i, step := range prompt {
		messages = append(messages, Message{
			Role:    "user",
			Content: step,
		})
		r, err := request(model, key, messages)
		if err != nil {
			return nil, 0, 0, err
		}
		inToken += r.Usage.InputTokens
		outToken += r.Usage.OutputTokens
		if len(r.Choices) == 0 {
			fmt.Printf("AI response err: %+v\n", r)
			return nil, 0, 0, err
		} else {
			messages = append(messages, r.Choices[0].Message)
		}
		fmt.Printf("AI 回答 %v: %v\n", i, r.Choices[0].Message.Content)
		if err = cb(r.Choices[0].Message.Content, i); err != nil {
			return nil, 0, 0, err
		}
	}
	fmt.Printf("AI finished, use token: in-%v out-%v\n", inToken, outToken)
	return messages, inToken, outToken, nil
}