package openai

import (
	"context"

	gogpt "github.com/sashabaranov/go-gpt3"

	"github.com/tidbcloud/aitestsuit/internal/entity"
)

var (
	c *gogpt.Client
)

func InitGPT(token string) {
	c = gogpt.NewClient(token)
}

func buildRequest(prompt string, params entity.Params) gogpt.CompletionRequest {
	req := gogpt.CompletionRequest{
		Model:       params.ModelType,
		Prompt:      prompt,
		MaxTokens:   params.MaxTokens,
		Temperature: params.Temperature,
		TopP:        params.TopP,
		N:           params.N,
		Stream:      false,
		Stop:        params.Stop,
	}

	return req
}

func GetReplyFromAI(ctx context.Context, promptText string, params entity.Params) (string, int, error) {
	req := buildRequest(promptText, params)
	resp, err := c.CreateCompletion(ctx, req)
	if err != nil {
		return "", resp.Usage.TotalTokens, err
	}

	line := resp.Choices[0].Text
	return line, resp.Usage.TotalTokens, nil
}