package chatrecords

import (
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"github.com/suyuan32/simple-admin-core/rpc/types/core"
	"github.com/zeromicro/go-zero/core/logx"
	"io"
	"net/http"
	"strconv"
	"strings"
	"time"
	"wechat-api/ent"
	"wechat-api/ent/creditbalance"
	"wechat-api/ent/custom_types"
	"wechat-api/ent/employee"
	"wechat-api/hook/credit"
	"wechat-api/hook/dify"
	"wechat-api/internal/svc"
	"wechat-api/internal/types"
	"wechat-api/internal/utils"
	jwtutils "wechat-api/internal/utils/jwt"
)

type GptsSubmitApiChatLogic struct {
	logx.Logger
	ctx    context.Context
	svcCtx *svc.ServiceContext
}

func NewGptsSubmitApiChatLogic(ctx context.Context, svcCtx *svc.ServiceContext) *GptsSubmitApiChatLogic {
	return &GptsSubmitApiChatLogic{
		Logger: logx.WithContext(ctx),
		ctx:    ctx,
		svcCtx: svcCtx}
}

func (l *GptsSubmitApiChatLogic) GptsSubmitApiChat(tokenStr string, req *types.GptChatReq, w http.ResponseWriter) {
	flusher, ok := w.(http.Flusher)
	if !ok {
		return
	}

	claims, err := jwtutils.ParseJwtToken(l.svcCtx.Config.Auth.AccessSecret, tokenStr)
	if err != nil {
		jsonData := ChatMessage{}
		jsonData.NeedLogin = true
		jsonDataStr, _ := json.Marshal(jsonData)
		_, _ = fmt.Fprintf(w, "%s", "data: "+string(jsonDataStr)+"\r\n")
		flusher.Flush()
		return
	}

	//userId, ok := "0190ee30-2b0c-739e-88e2-c59d4d1175c1", true
	userId, ok := claims["userId"].(string) //用户的uuid
	if !ok || userId == "" {
		jsonData := ChatMessage{}
		jsonData.NeedLogin = true
		jsonDataStr, _ := json.Marshal(jsonData)
		_, _ = fmt.Fprintf(w, "%s", "data: "+string(jsonDataStr)+"\r\n")
		flusher.Flush()
		return
	}
	userInfo, _ := l.svcCtx.CoreRpc.GetUserById(l.ctx, &core.UUIDReq{Id: userId})
	valid := utils.CheckGptLogin(userInfo.RoleIds)
	//l.Logger.Error(valid)
	if !valid {
		jsonData := ChatMessage{}
		jsonData.NeedLogin = true
		jsonDataStr, _ := json.Marshal(jsonData)
		_, _ = fmt.Fprintf(w, "%s", "data: "+string(jsonDataStr)+"\r\n")
		flusher.Flush()
		return
	}

	// 查询积分余额 是否足够使用
	// 先查用户的积分余额，如果不足 则查租户的积分余额
	var switcher bool
	var balance float64
	switcher = false //TODO 打开这里就是开始计费，并且会记录使用积分和token的信息
	creditBalance, err := l.svcCtx.DB.CreditBalance.Query().Where(creditbalance.OrganizationID(*userInfo.DepartmentId)).Only(l.ctx)
	if switcher {
		if err != nil {
			if !ent.IsNotFound(err) {
				l.Logger.Errorf("query user balance failed:%v\n", err)
				return
			}
		} else {
			if creditBalance.Balance > 0 {
				balance = creditBalance.Balance
			}
		}

		//l.Logger.Error(balance)
		if balance <= 0 {
			jsonData := ChatMessage{}
			jsonData.NeedPay = true
			jsonDataStr, _ := json.Marshal(jsonData)
			_, _ = fmt.Fprintf(w, "%s", "data: "+string(jsonDataStr)+"\r\n")
			flusher.Flush()
			return
		}
	}

	employeeItem, err := l.svcCtx.DB.Employee.Query().Where(employee.ID(*req.AgentId)).Only(l.ctx)
	if err != nil {
		if ent.IsNotFound(err) {
			l.Logger.Error("employee not found")
			return
		}
		l.Logger.Errorf("query employee failed:%v\n", err)
		return
	}

	var chatReq dify.ChatReq
	chatReq.ResponseMode = "streaming"
	chatReq.Query = *req.Content
	chatReq.User = fmt.Sprintf("%s-%d", userId, *req.AgentId)

	chatReq.ConversationId = *req.ConversationId

	// 設置請求體 (這裡是一個簡單的範例，你可以根據需要調整)
	jsonData, _ := json.Marshal(chatReq)
	fmt.Printf("request data:%v\n", string(jsonData))

	// 建立HTTP請求
	url := employeeItem.APIBase + dify.GetChatUrl()
	request, err := http.NewRequest("POST", url, bytes.NewBuffer(jsonData))
	if err != nil {
		fmt.Printf("Error creating request:%v", err)
		return
	}

	// 設置請求頭
	request.Header.Set("Content-Type", "application/json")
	request.Header.Set("Authorization", "Bearer "+employeeItem.APIKey)
	request.Header.Set("Transfer-Encoding", "chunked")

	// 發送請求
	client := &http.Client{
		Timeout: time.Second * 60,
	}
	response, err := client.Do(request)
	defer response.Body.Close()

	// 讀取響應
	reader := bufio.NewReader(response.Body)

	w.Header().Set("Content-Type", "text/event-stream;charset=utf-8")
	w.Header().Set("Connection", "keep-alive")
	w.Header().Set("Cache-Control", "no-cache")

	var answer string
	var finish bool
	for {
		line, err := reader.ReadString('\n')
		line = strings.Trim(line, " \n")

		if err == io.EOF {
			break
		}

		if len(line) > 6 {
			line = line[6:]
			chatData := dify.ChatResp{}
			err = json.Unmarshal([]byte(line), &chatData)
			if err != nil {
				fmt.Printf("json unmarshall error:%v\n", err)
				fmt.Printf("line:%v\n", line)
				continue
			}

			answer = answer + chatData.Answer

			if chatData.Event == "message_end" {
				finish = true
			}

			jsonData := ChatMessage{}
			jsonData.Id = chatData.Id
			jsonData.Answer = chatData.Answer
			jsonData.MessageId = chatData.MessageId
			jsonData.Finish = finish
			jsonData.ConversationId = chatData.ConversationId
			lineData, _ := json.Marshal(jsonData)

			_, err = fmt.Fprintf(w, "%s", "data: "+string(lineData)+"\r\n")
			fmt.Printf("response=%v\n", string(lineData))
			if err != nil {
				fmt.Printf("Error writing to client:%v \n", err)
				break
			}
			flusher.Flush()

			if finish {
				if switcher {
					tx, err := l.svcCtx.DB.Tx(context.Background())
					agentId := strconv.Itoa(int(*req.AgentId))

					// 构造 original_data
					originalData := custom_types.OriginalData{}
					originalData.Request = chatReq
					originalData.Response = chatData

					if err != nil {
						l.Logger.Errorf("start transaction error:%v\n", err)
					} else {
						err = credit.AddCreditUsage(tx, l.ctx,
							agentId, userId, *userInfo.DepartmentId,
							req.Content, &answer,
							&originalData, &chatData,
						)
						if err != nil {
							_ = tx.Rollback()
							l.Logger.Errorf("save credits info failed:%v\n", err)
						} else {
							_ = tx.Commit()
						}
					}
				}
				break
			}
		}
	}
}
