package middleware

import (
	"bytes"
	"io"

	"github.com/gin-gonic/gin"
	log "github.com/sirupsen/logrus"
)

func PrintRequestResponseBodyMiddleware() gin.HandlerFunc {
	return func(c *gin.Context) {
		// 读取请求 body
		var requestBody []byte
		if c.Request.Body != nil {
			requestBody, _ = io.ReadAll(c.Request.Body)
			c.Request.Body = io.NopCloser(bytes.NewBuffer(requestBody))
		}

		log.WithFields(log.Fields{"method": c.Request.Method, "uri": c.Request.RequestURI, "body": string(requestBody)}).Debug("request body")

		bodyWriter := &responseBodyWriter{body: bytes.NewBufferString(""), ResponseWriter: c.Writer}
		c.Writer = bodyWriter

		c.Next()

		responseBody := bodyWriter.body.String()
		log.WithFields(log.Fields{"status": c.Writer.Status(), "body": responseBody}).Debug("response body")
	}
}

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)
}
