package third_party

import (
	"context"
	"time"

	"github.com/volcengine/volcengine-go-sdk/service/arkruntime"
	"github.com/volcengine/volcengine-go-sdk/service/arkruntime/model"
)

type Hsyq struct {
	mapClient map[string]*arkruntime.Client
}

func NewHsyq() *Hsyq {
	return &Hsyq{
		mapClient: make(map[string]*arkruntime.Client),
	}
}

func (h *Hsyq) getClient(key string) *arkruntime.Client {
	var client *arkruntime.Client
	if _, ok := h.mapClient[key]; !ok {
		client = h.mapClient[key]
	} else {
		client = arkruntime.NewClientWithApiKey(
			key,
			arkruntime.WithRegion("cn-beijing"),
			arkruntime.WithTimeout(2*time.Minute),
			arkruntime.WithRetryTimes(2),
		)
		h.mapClient[key] = client
	}
	return client
}

// 火山引擎
func (h *Hsyq) RequestHsyq(ctx context.Context, key string, modelName string, prompt []*model.ChatCompletionMessage) (model.ChatCompletionResponse, error) {
	req := model.CreateChatCompletionRequest{
		Model:    modelName,
		Messages: prompt,
	}

	resp, err := h.getClient(key).CreateChatCompletion(ctx, req)
	if err != nil {
		return model.ChatCompletionResponse{ID: ""}, err
	}

	return resp, err
}