package chatrecords

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

	"wechat-api/internal/svc"
	"wechat-api/internal/types"

	"github.com/zeromicro/go-zero/core/logx"
)

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 := 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)
	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 float32
	//switcher = true	//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

			}
		}

		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())
					// 构造 original_data
					originalData := custom_types.OriginalData{}
					originalData.Request = chatReq
					originalData.Response = chatData

					agentId := strconv.Itoa(int(*req.AgentId))
					// 记录Token使用信息
					usageDetailItem, err := tx.UsageDetail.Create().
						SetType(3).            //1-微信 2-名片 3-智能体
						SetBotID(agentId).     //智能体ID
						SetReceiverID(userId). //接收者userID
						SetApp(8). //8-智能体
						SetSessionID(0).
						SetRequest(*req.Content).
						SetResponse(answer).
						SetOriginalData(originalData).
						SetTotalTokens(chatData.Metadata.Usage.TotalTokens).
						SetPromptTokens(chatData.Metadata.Usage.PromptTokens).
						SetCompletionTokens(chatData.Metadata.Usage.CompletionTokens).
						Save(l.ctx)

					if err != nil {
						_ = tx.Rollback()
						l.Logger.Errorf("save data to usage_detail error:%v\n", err)
					}

					// 积分明细表记录使用量
					// 根据1:1000 根据Token换算积分使用量
					change := float64(chatData.Metadata.Usage.TotalTokens) / float64(1000)
					number := float32(math.Round(change*1000) / 1000)
					_, err = tx.CreditUsage.Create().
						SetUserID(userId).
						SetNumber(number).
						SetNtype(1).
						SetNid(usageDetailItem.ID).
						SetTable("usage_detail").
						SetOrganizationID(*userInfo.DepartmentId).
						Save(l.ctx)
					if err != nil {
						_ = tx.Rollback()
						l.Logger.Errorf("save data to credit_usage error:%v\n", err)
					}

					// 减积分
					_, err = tx.CreditBalance.Update().Where(creditbalance.OrganizationID(*userInfo.DepartmentId)).SetBalance(creditBalance.Balance - number).Save(l.ctx)
					if err != nil {
						_ = tx.Rollback()
						l.Logger.Errorf("update organization:%v balance error:%v\n", *userInfo.DepartmentId, err)
					}
					_ = tx.Commit()
				}
				break
			}
		}
	}
}
