package v1

import (
	"bytes"
	"context"
	"encoding/json"
	"io"
	"net/http"
	"strings"
	"time"

	"github.com/pingcap-inc/pingcat/internal/pkg/code"
	v1 "github.com/pingcap-inc/pingcat/model/v1"
	"github.com/spf13/viper"

	"github.com/pingcap-inc/pchub/pkg/errors"
	"github.com/pingcap-inc/pchub/pkg/log"
)

const (
	ChatURL = "https://api.openai.com/v1/chat/completions"
)

type OpenAIChatReqData struct {
	Model       string               `json:"model,omitempty"`
	Temperature float32              `json:"temperature,omitempty"`
	TopP        float32              `json:"top_p,omitempty"`
	MaxTokens   int                  `json:"max_tokens,omitempty"`
	Messages    []*v1.GPTChatMessage `json:"messages"`
}

type OpenAIChatResData struct {
	Model   string                    `json:"model"`
	Choices []OpenAIChatResDataChoice `json:"choices"`
	Usage   OpenAIChatResDataUsage    `json:"usage"`
}

type OpenAIChatResDataChoice struct {
	Message *v1.GPTChatMessage `json:"message"`
}

type OpenAIChatResDataUsage struct {
	PromptTokens     int `json:"prompt_tokens"`
	CompletionTokens int `json:"completion_tokens"`
	TotalTokens      int `json:"total_tokens"`
}

type OpenAIChatErrRes struct {
	Error OpenAIChatErrResError `json:"error"`
}

type OpenAIChatErrResError struct {
	Code    string `json:"code"`
	Message string `json:"message"`
}

const (
	OAIErrCodeLengthExceeded = "context_length_exceeded"
)

type ReqSetting struct {
	Model       string
	Temperature float32
	TopP        float32
	MaxTokens   int
}

func (s *GPTService) constructReqBody(ctx context.Context, msgs []*v1.GPTChatMessage, reqSetting *ReqSetting) ([]byte, error) {
	reqBody := OpenAIChatReqData{
		Messages: append([]*v1.GPTChatMessage{
			{
				Role:    "system",
				Content: "You are an in-house assistant to PingCAP, that means you serve as an internal member of the PingCAP company. Your name is PingCAT and your avatar is a cat image created by Midjourney(another image AI).",
			},
		}, msgs...),
	}
	//if reqSetting.Model != "gpt4" {
	//	reqBody.Model = "gpt-3.5-turbo"
	//}

	if reqSetting.Temperature != 0.0 {
		reqBody.Temperature = reqSetting.Temperature
	}
	if reqSetting.TopP != 0.0 {
		reqBody.TopP = reqSetting.TopP
	}
	if reqSetting.MaxTokens != 0 {
		reqBody.MaxTokens = reqSetting.MaxTokens
	}
	reqb, err := json.Marshal(&reqBody)
	if err != nil {
		log.L(ctx).Error(err.Error())

		return nil, errors.WithCode(code.ErrEncodingJSON, err.Error())
	}

	return reqb, nil
}

func (s *GPTService) constructSetting(ctx context.Context, larkUserID string, opts *v1.GPTChatOption) *ReqSetting {
	log.L(ctx).Info("GPTService.constructSetting function called.")

	settings := &ReqSetting{
		Model:       "gpt3",
		Temperature: 0.0,
		TopP:        0.0,
		MaxTokens:   0,
	}
	uSettings, err := s.store.UserOpenAI().GetByUser(ctx, larkUserID)
	if err == nil && uSettings.ID > 0 {
		if uSettings.Model == "gpt4" {
			settings.Model = "gpt4"
		}
		if uSettings.Temperature != 0.0 {
			settings.Temperature = uSettings.Temperature
		}
		if uSettings.TopP != 0.0 {
			settings.TopP = uSettings.TopP
		}
		if uSettings.MaxTokens != 0 {
			settings.MaxTokens = uSettings.MaxTokens
		}
	}
	if opts != nil {
		if opts.Temperature != 0.0 {
			settings.Temperature = opts.Temperature
		}
		if opts.TopP != 0.0 {
			settings.TopP = opts.TopP
		}
		if opts.MaxTokens != 0 {
			settings.MaxTokens = opts.MaxTokens
		}
	}

	return settings
}

func (s *GPTService) constructReq(ctx context.Context, msgs []*v1.GPTChatMessage, opts *v1.GPTChatOption) (*http.Request, error) {
	log.L(ctx).Info("GPTService.constructReq function called.")

	larkUserID, ok := ctx.Value("larkUserID").(string)
	if !ok || larkUserID == "" {
		return nil, errors.WithCode(code.ErrNoUserLarkID, "No user lark id in context.")
	}
	settings := s.constructSetting(ctx, larkUserID, opts)
	reqb, err := s.constructReqBody(ctx, msgs, settings)
	if err != nil {
		return nil, err
	}
	url := ChatURL
	if settings.Model == "gpt4" {
		url = viper.GetString("azure.gpt4Endpoint")
	} else {
		url = viper.GetString("azure.gpt3Endpoint")
	}
	log.L(ctx).Infof("Will call gpt api, endpoint: %s, req data: %s", url, string(reqb))
	req, _ := http.NewRequestWithContext(ctx, "POST", url, bytes.NewReader(reqb))
	req.Header.Set("Content-Type", "application/json")
	if settings.Model == "gpt4" {
		req.Header.Set("api-key", viper.GetString("azure.gpt4Key"))
		q := req.URL.Query()
		q.Add("api-version", viper.GetString("azure.gpt4Version"))
		req.URL.RawQuery = q.Encode()
	} else {
		req.Header.Set("api-key", viper.GetString("azure.gpt3Key"))
		q := req.URL.Query()
		q.Add("api-version", viper.GetString("azure.gpt3Version"))
		req.URL.RawQuery = q.Encode()
		//req.Header.Set("Authorization", "Bearer "+viper.GetString("openai.key"))
		//req.Header.Set("OpenAI-Organization", viper.GetString("openai.orgId"))
	}

	return req, nil
}

func (s *GPTService) ChatCompletion(ctx context.Context, msgs []*v1.GPTChatMessage, opts *v1.GPTChatOption) (*v1.GPTChatRes, error) {
	log.L(ctx).Info("GPTService.ChatCompletion function called.")

	req, err := s.constructReq(ctx, msgs, opts)
	if err != nil {
		return nil, err
	}

	tStartNano := time.Now().UnixNano()
	resp, err := (&http.Client{}).Do(req)
	tEndNano := time.Now().UnixNano()
	log.L(ctx).Info("Call openai cost: " + time.Duration(tEndNano-tStartNano).String())
	if err != nil {
		log.L(ctx).Error(err.Error())

		return nil, errors.WithCode(code.ErrOpenAIHttpRequest, "error when request openai chat api, err: %s", err.Error())
	}
	defer func() {
		_ = resp.Body.Close()
	}()

	b, _ := io.ReadAll(resp.Body)

	if resp.StatusCode == http.StatusOK {
		var res OpenAIChatResData
		if err := json.Unmarshal(b, &res); err != nil {
			log.L(ctx).Error(err.Error())

			return nil, errors.WithCode(code.ErrDecodingJSON, "error when unmarshal openai chat api res, err: %s", err.Error())
		}

		if len(res.Choices) == 0 || res.Choices[0].Message == nil {
			return nil, errors.WithCode(code.ErrNoValidMessage, "no valid message content in openai response, err: %s", err.Error())
		}

		return &v1.GPTChatRes{
			Content: strings.TrimSpace(res.Choices[0].Message.Content),
			Role:    res.Choices[0].Message.Role,
			Usage:   v1.GPTTokenUsage{Prompt: res.Usage.PromptTokens, Completion: res.Usage.CompletionTokens, Total: res.Usage.TotalTokens},
			Model:   res.Model,
		}, nil
	}
	var res OpenAIChatErrRes
	if err := json.Unmarshal(b, &res); err != nil {
		log.L(ctx).Error(err.Error())

		return nil, errors.WithCode(code.ErrDecodingJSON, "error when unmarshal openai chat api err res, err: %s", err.Error())
	}
	if res.Error.Code == OAIErrCodeLengthExceeded {
		return nil, errors.WithCode(code.ErrContextLengthExceeded, "error when request openai chat api, err: %s", res.Error.Message)
	}

	return nil, errors.WithCode(code.ErrOpenAIHttpRequest, "error when request openai chat api, err: %s", res.Error.Message)
}

const PingCAPORGLarkUserID = "pingcat"