package v1

import (
	"context"

	"github.com/pingcap-inc/pingcat/internal/pkg/code"
	v1 "github.com/pingcap-inc/pingcat/model/v1"
	uuid "github.com/satori/go.uuid"

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

func (s *ChatService) ReplyGeneralMsg(ctx context.Context, d v1.MsgReqData) (*v1.MessageResData, error) {
	log.L(ctx).Info("ChatSrv.ReplyGeneralMsg called")

	msgID := uuid.Must(uuid.NewV4(), nil).String()
	inMsg, thmsgs := s.constructInMsg(ctx, d, v1.SceneGeneral, msgID)
	inMsgDB, err := s.store.Message().Create(ctx, inMsg)
	if err != nil {
		log.L(ctx).Error(errors.Wrap(err, "Error when input create msg in db").Error())

		return nil, errors.WithCode(code.ErrDatabase, "")
	}
	outMsgCont := ""
	var usage *v1.GPTTokenUsage
	gptMsgs := s.convert2GPTMsg(ctx, thmsgs, inMsg)
	gptSrv := NewGPT(s.store)
	gptMsg, err := gptSrv.ChatCompletion(ctx, gptMsgs, d.Options)
	if err != nil {
		coder := errors.ParseCoder(err)
		if coder.Code() != code.ErrContextLengthExceeded {
			return nil, errors.Wrap(err, "error when request openai api.")
		}
		outMsgCont = "Context length exceeded, please start a new thread with no 'pid' in the request body. Err: " + err.Error()
	} else {
		outMsgCont = gptMsg.Content
		usage = &gptMsg.Usage
	}
	outMsgID := uuid.Must(uuid.NewV4(), nil).String()
	outMsg := s.constructOutMsg(ctx, usage, outMsgCont, v1.SceneGeneral, outMsgID, inMsgDB, "", gptMsg.Model)
	outMsgDB, err := s.store.Message().Create(ctx, outMsg)
	if err != nil {
		log.L(ctx).Error(errors.Wrap(err, "Error when output create msg in db").Error())

		return nil, errors.WithCode(code.ErrDatabase, "")
	}

	return &v1.MessageResData{
		InMsg:  *inMsgDB,
		OutMsg: *outMsgDB,
		Extra: v1.MessageResDataExtra{
			GPTTokenUsage: &(gptMsg.Usage),
			GPTOptions:    d.Options,
		},
	}, nil
}

func (s *ChatService) convert2GPTMsg(_ context.Context, pmsgs []*v1.Message, msg *v1.Message) []*v1.GPTChatMessage {
	gptMsgs := make([]*v1.GPTChatMessage, 0)
	for _, it := range pmsgs {
		cur := &v1.GPTChatMessage{
			Content: it.Content,
		}
		if it.SenderType == v1.SubjectTypeUser {
			cur.Role = "user"
		} else {
			cur.Role = "assistant"
		}
		gptMsgs = append(gptMsgs, cur)
	}

	gptMsgs = append(gptMsgs, &v1.GPTChatMessage{
		Content: msg.Content,
		Role:    "user",
	})

	return gptMsgs
}

func (s *ChatService) constructOutMsg(ctx context.Context, usage *v1.GPTTokenUsage, outMsgCont, scene, id string, inMsgDB *v1.Message, groupID, model string) *v1.Message {
	larkUserID, _ := ctx.Value("larkUserID").(string)
	outMsg := &v1.Message{
		ID:             id,
		Scene:          scene,
		Content:        outMsgCont,
		SenderLarkID:   PingCAPORGLarkUserID,
		SenderType:     v1.SubjectTypePingCat,
		ReceiverType:   v1.SubjectTypeUser,
		ReceiverLarkID: larkUserID,
		MsgType:        "text",
		PID:            inMsgDB.ID,
		Model:          model,
	}
	if usage != nil {
		outMsg.TotalTokens = usage.Total
		outMsg.PromptTokens = usage.Prompt
		outMsg.CompletionTokens = usage.Completion
	}
	if groupID != "" {
		outMsg.ReceiverLarkID = groupID
		outMsg.ReceiverType = v1.SubjectTypeGroup
	}
	if inMsgDB.FlatParentIDs != "" {
		outMsg.FlatParentIDs = inMsgDB.FlatParentIDs + "," + inMsgDB.ID
	} else {
		outMsg.FlatParentIDs = inMsgDB.ID
	}

	return outMsg
}

func (s *ChatService) constructInMsg(ctx context.Context, d v1.MsgReqData, scene, id string) (*v1.Message, []*v1.Message) {
	larkUserID, _ := ctx.Value("larkUserID").(string)
	inMsg := &v1.Message{
		ID:             id,
		Scene:          scene,
		Content:        d.Msg,
		SenderLarkID:   larkUserID,
		SenderType:     v1.SubjectTypeUser,
		ReceiverLarkID: PingCAPORGLarkUserID,
		ReceiverType:   v1.SubjectTypePingCat,
		MsgType:        "text",
	}

	thmsgs := make([]*v1.Message, 0)
	if d.PID != "" {
		msgs, err := s.listThreadMsgs(ctx, d.PID)
		if err == nil && len(msgs) > 0 {
			thmsgs = msgs
		}
	}
	if len(thmsgs) > 0 {
		inMsg.PID = d.PID
		inMsg.FlatParentIDs = thmsgs[len(thmsgs)-1].FlatParentIDs
		if inMsg.FlatParentIDs == "" {
			inMsg.FlatParentIDs = d.PID
		} else {
			inMsg.FlatParentIDs = inMsg.FlatParentIDs + "," + d.PID
		}
	} else {
		inMsg.PID = ""
		inMsg.FlatParentIDs = ""
	}

	return inMsg, thmsgs
}

func (s *ChatService) listThreadMsgs(ctx context.Context, pid string) ([]*v1.Message, error) {
	pmsgs, err := s.store.Message().ListThreadMsgs(ctx, pid)
	if err != nil {
		log.L(ctx).Error("error when retrieve parent/ancestor messages for pid: " + pid)

		return nil, err
	}

	return pmsgs, nil
}