package middleware

import (
	"io"
	"net/http"

	"yh_web/server/pkg/traffic"

	"github.com/gin-gonic/gin"
)

type countReadCloser struct {
	io.ReadCloser
}

func (c *countReadCloser) Read(p []byte) (int, error) {
	n, err := c.ReadCloser.Read(p)
	if n > 0 {
		traffic.AddIn(n)
	}
	return n, err
}

type meterResponseWriter struct {
	gin.ResponseWriter
}

func (w *meterResponseWriter) Write(p []byte) (int, error) {
	n, err := w.ResponseWriter.Write(p)
	if n > 0 {
		traffic.AddOut(n)
	}
	return n, err
}

func (w *meterResponseWriter) WriteString(s string) (int, error) {
	n, err := w.ResponseWriter.WriteString(s)
	if n > 0 {
		traffic.AddOut(n)
	}
	return n, err
}

// TrafficMeter 统计 HTTP 请求体与响应体字节量（进程级，非网卡级）。
func TrafficMeter() gin.HandlerFunc {
	return func(c *gin.Context) {
		if c.Request.Body != nil && c.Request.Body != http.NoBody {
			c.Request.Body = &countReadCloser{ReadCloser: c.Request.Body}
		}
		c.Writer = &meterResponseWriter{ResponseWriter: c.Writer}
		c.Next()
	}
}