package doubao

import (
	"context"
	"github.com/volcengine/volcengine-go-sdk/service/arkruntime"
	"github.com/volcengine/volcengine-go-sdk/service/arkruntime/model"
	"time"
)

type DouBao struct {
	Model string
	Key   string
}

func NewDouBao(modelName string, key string) *DouBao {
	return &DouBao{
		Model: modelName,
		Key:   key,
	}
}

func (o *DouBao) GetData(ctx context.Context, respHandle func(input string) (string, error), Message []*model.ChatCompletionMessage) (string, error) {

	client := arkruntime.NewClientWithApiKey(
		o.Key,
		//arkruntime.WithBaseUrl(UrlMap[url]),
		arkruntime.WithRegion("cn-beijing"),
		arkruntime.WithTimeout(2*time.Minute),
		arkruntime.WithRetryTimes(2),
	)
	req := model.CreateChatCompletionRequest{
		Model:    o.Model,
		Messages: Message,
	}

	resp, err := client.CreateChatCompletion(ctx, req)
	if err != nil {
		return "", err
	}
	result, err := respHandle(*resp.Choices[0].Message.Content.StringValue)

	return result, err
}