package middleware

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

	"git.echol.cn/loser/ai_proxy/server/global"
	"git.echol.cn/loser/ai_proxy/server/model/system"
	"git.echol.cn/loser/ai_proxy/server/service"
	"git.echol.cn/loser/ai_proxy/server/utils"
	"github.com/gin-gonic/gin"
	"go.uber.org/zap"
)

var operationRecordService = service.ServiceGroupApp.SystemServiceGroup.OperationRecordService

// OperationRecord 操作记录中间件
func OperationRecord() gin.HandlerFunc {
	return func(c *gin.Context) {
		var body []byte
		var userId int
		if c.Request.Method != http.MethodGet {
			var err error
			body, err = io.ReadAll(c.Request.Body)
			if err != nil {
				global.GVA_LOG.Error("read body from request error:", zap.Error(err))
			} else {
				c.Request.Body = io.NopCloser(bytes.NewBuffer(body))
			}
		}

		userId = int(utils.GetUserID(c))

		writer := responseBodyWriter{
			ResponseWriter: c.Writer,
			body:           &bytes.Buffer{},
		}
		c.Writer = writer
		now := time.Now()

		c.Next()

		latency := time.Since(now)

		if c.Request.Method != http.MethodGet {
			record := system.SysOperationRecord{
				Ip:      c.ClientIP(),
				Method:  c.Request.Method,
				Path:    c.Request.URL.Path,
				Agent:   c.Request.UserAgent(),
				Body:    string(body),
				UserID:  userId,
				Status:  c.Writer.Status(),
				Latency: latency,
				Resp:    writer.body.String(),
			}

			values, _ := url.ParseQuery(c.Request.URL.RawQuery)
			record.Query = values.Encode()

			if err := operationRecordService.CreateSysOperationRecord(record); err != nil {
				global.GVA_LOG.Error("create operation record error:", zap.Error(err))
			}
		}
	}
}

type responseBodyWriter struct {
	gin.ResponseWriter
	body *bytes.Buffer
}

func (r responseBodyWriter) Write(b []byte) (int, error) {
	r.body.Write(b)
	return r.ResponseWriter.Write(b)
}

func (r responseBodyWriter) WriteString(s string) (int, error) {
	r.body.WriteString(s)
	return r.ResponseWriter.WriteString(s)
}

func (r responseBodyWriter) WriteJSON(obj interface{}) error {
	data, err := json.Marshal(obj)
	if err != nil {
		return err
	}
	r.body.Write(data)
	return r.ResponseWriter.WriteJSON(obj)
}

// NeedRecordPath 判断是否需要记录操作日志
func NeedRecordPath(path string) bool {
	// 排除不需要记录的路径
	excludePaths := []string{
		"/health",
		"/swagger",
		"/api/captcha",
	}

	for _, excludePath := range excludePaths {
		if strings.HasPrefix(path, excludePath) {
			return false
		}
	}

	return true
}