package middleware

import (
	"bytes"
	"io"
	"strings"
	"time"

	"gitea.com/bitwsd/document_ai/pkg/log"
	"github.com/gin-gonic/gin"
)

const (
	maxBodySize = 1024 * 500 // 500KB 限制
)

// 自定义 ResponseWriter 来捕获响应
type bodyWriter struct {
	gin.ResponseWriter
	body *bytes.Buffer
}

func (w *bodyWriter) Write(b []byte) (int, error) {
	w.body.Write(b)
	return w.ResponseWriter.Write(b)
}

func AccessLog() gin.HandlerFunc {
	return func(c *gin.Context) {
		start := time.Now()
		path := c.Request.URL.Path
		raw := c.Request.URL.RawQuery

		// 处理请求体
		var reqBody string
		if c.Request.Body != nil && (c.Request.Method == "POST" || c.Request.Method == "PUT") {
			// 读取并限制请求体大小
			bodyBytes, _ := io.ReadAll(io.LimitReader(c.Request.Body, maxBodySize))
			c.Request.Body = io.NopCloser(bytes.NewBuffer(bodyBytes)) // 重新设置 body
			reqBody = string(bodyBytes)
		}

		// 设置自定义 ResponseWriter
		var responseBody string
		if !strings.Contains(c.GetHeader("Accept"), "text/event-stream") {
			bw := &bodyWriter{body: &bytes.Buffer{}, ResponseWriter: c.Writer}
			c.Writer = bw
		}

		c.Next()

		// 获取响应体（非 SSE）
		if writer, ok := c.Writer.(*bodyWriter); ok {
			responseBody = writer.body.String()
			if len(responseBody) > maxBodySize {
				responseBody = responseBody[:maxBodySize] + "... (truncated)"
			}
		}

		// 记录访问日志
		log.Access(c.Request.Context(),
			"request_id", c.GetString("request_id"),
			"method", c.Request.Method,
			"path", path,
			"query", raw,
			"ip", c.ClientIP(),
			"user_agent", c.Request.UserAgent(),
			"status", c.Writer.Status(),
			"duration", time.Since(start),
			"request_body", reqBody,
			"response_body", responseBody,
		)
	}
}