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, messages ...Message) (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{},
		Stream:    true,
	}
	for _, m := range messages {
		req.Messages = append(req.Messages, openai.ChatCompletionMessage{
			Role:    m.Role,
			Content: m.Content,
		})
	}
	req.Messages = append(req.Messages, openai.ChatCompletionMessage{
		Role:    openai.ChatMessageRoleUser,
		Content: question,
	})
	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
}