package ai

import (
	"context"
	"errors"
	"fmt"
	"github.com/sashabaranov/go-openai"
	"github.com/zeromicro/go-zero/core/logx"
	"io"
)

func ChatGPT(gptModelCode, key string, question string, channel chan string) (answer string, err error) {
	config := openai.DefaultConfig(key)
	config.BaseURL = "http://47.251.84.160:8080/v1" // "https://api.openai.com/v1" //
	c := openai.NewClientWithConfig(config)
	ctx := context.Background()

	req := openai.ChatCompletionRequest{
		Model:     gptModelCode,
		MaxTokens: 2048,
		Messages: []openai.ChatCompletionMessage{
			{
				Role:    openai.ChatMessageRoleUser,
				Content: question,
			},
		},
		Stream: true,
	}
	stream, err := c.CreateChatCompletionStream(ctx, req)
	if err != nil {
		fmt.Printf("ChatCompletionStream error: %v\n", err)
		return
	}
	defer func() {
		stream.Close()
		return
	}()

	logx.Info("Stream response: ")
	for {
		response, recvErr := stream.Recv()
		if errors.Is(recvErr, io.EOF) {
			fmt.Println("\nStream finished")
			break
		}
		logx.Info(response)
		if recvErr != nil {
			logx.Info("\nStream error: %v\n", recvErr)
			break
		}
		content := response.Choices[0].Delta.Content
		answer += content
		channel <- content
	}
	return
}