package doubao

import (
	"context"
	"github.com/volcengine/volcengine-go-sdk/service/arkruntime"
	"github.com/volcengine/volcengine-go-sdk/service/arkruntime/model"
	"github.com/volcengine/volcengine-go-sdk/volcengine"
	"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, url UrlType, respHandle func(input string) (string, error), text ...string) (string, error) {
	var Message = make([]*model.ChatCompletionMessage, len(text))

	key := "914ccf1d-c002-4fad-a431-f291f5e0d2ad"
	client := arkruntime.NewClientWithApiKey(
		key,
		//arkruntime.WithBaseUrl(UrlMap[url]),
		arkruntime.WithRegion("cn-beijing"),
		arkruntime.WithTimeout(2*time.Minute),
		arkruntime.WithRetryTimes(2),
	)
	for k, v := range text {
		Message[k] = &model.ChatCompletionMessage{
			Role: model.ChatMessageRoleSystem,
			Content: &model.ChatCompletionMessageContent{
				StringValue: volcengine.String(v),
			},
		}
	}
	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
}