package channel

import (
	"bytes"
	"context"
	"github.com/openai/openai-go"
	"github.com/openai/openai-go/option"
	"github.com/zeromicro/go-zero/core/logx"
	"io"
	"io/ioutil"
	"net/http"
	"strings"
	"time"
)

type ChatEngine struct {
	BaseURL   string
	APIKey    string
	IsFastGpt bool
}

func NewChatEngine(baseUrl string, apiKey string, isFastGpt bool) *ChatEngine {

	if baseUrl == "" {
		baseUrl = "https://api.openai.com/v1/"
	}

	if !strings.HasSuffix(baseUrl, "/") {
		baseUrl += "/"
	}

	return &ChatEngine{
		BaseURL:   baseUrl,
		APIKey:    apiKey,
		IsFastGpt: isFastGpt,
	}
}

func (o *ChatEngine) ChatCompletions(ctx context.Context, model openai.ChatModel, msg string, chatId string) (string, error) {

	client := openai.NewClient(
		option.WithBaseURL(o.BaseURL),
		option.WithAPIKey(o.APIKey),
		option.WithMiddleware(o.Logger),
	)

	options := make([]option.RequestOption, 0)

	if o.IsFastGpt {
		options = append(options, option.WithJSONSet("chatId", chatId))
	}

	param := openai.ChatCompletionNewParams{
		Messages: openai.F([]openai.ChatCompletionMessageParamUnion{
			openai.UserMessage(msg),
		}),
		Model: openai.F(model),
	}

	chatCompletion, err := client.Chat.Completions.New(ctx, param, options...)
	if err != nil {
		return "", err
	}

	return chatCompletion.Choices[0].Message.Content, nil
}

func (o *ChatEngine) Logger(req *http.Request, next option.MiddlewareNext) (res *http.Response, err error) {
	// Before the request
	start := time.Now()
	//logx.Info("request:", req)
	// 读取并打印请求的 body
	var requestBody []byte
	if req.Body != nil {
		requestBody, err = ioutil.ReadAll(req.Body)
		if err != nil {
			logx.Error("读取请求 body 失败:", err)
			return nil, err
		}
		// 重新设置请求的 body
		req.Body = io.NopCloser(bytes.NewBuffer(requestBody))
	}

	logx.Info("request:", req)
	logx.Info("request body:", string(requestBody))

	// Forward the request to the next handler
	res, err = next(req)

	// Handle stuff after the request
	end := time.Now()

	//耗时

	logx.Info("resp: ", res, err, end.Sub(start))

	return res, err
}