package middleware

import (
	"compress/gzip"
	"github.com/gin-gonic/gin"
	"net/http"
	"strings"
)

func GzipMiddleware() gin.HandlerFunc {
	return func(c *gin.Context) {
		// 1. 处理请求解压
		if c.Request.Header.Get("Content-Encoding") == "gzip" {
			gzReader, err := gzip.NewReader(c.Request.Body)
			if err != nil {
				c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": "Invalid gzip format"})
				return
			}
			defer gzReader.Close()
			c.Request.Body = gzReader
		}

		// 2. 设置响应压缩支持
		if strings.Contains(c.Request.Header.Get("Accept-Encoding"), "gzip") {
			// 创建gzip writer
			gzWriter := gzip.NewWriter(c.Writer)
			defer gzWriter.Close()

			// 替换原始writer为压缩writer
			originalWriter := c.Writer
			c.Writer = &gzipResponseWriter{
				ResponseWriter: originalWriter,
				gzWriter:       gzWriter,
			}

			// 设置响应头
			c.Header("Content-Encoding", "gzip")
			c.Header("Vary", "Accept-Encoding")
		}

		c.Next()
	}
}

// 自定义ResponseWriter实现gzip压缩
type gzipResponseWriter struct {
	gin.ResponseWriter
	gzWriter *gzip.Writer
}

func (w *gzipResponseWriter) Write(data []byte) (int, error) {
	return w.gzWriter.Write(data)
}

func (w *gzipResponseWriter) WriteString(s string) (int, error) {
	return w.gzWriter.Write([]byte(s))
}