package middleware

import (
	"compress/gzip"
	"strings"
	"sync"

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

var gzPool = sync.Pool{
	New: func() any {
		w, _ := gzip.NewWriterLevel(nil, gzip.DefaultCompression)
		return w
	},
}

type gzResponseWriter struct {
	gin.ResponseWriter
	gz *gzip.Writer
}

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

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

func (w *gzResponseWriter) WriteHeader(code int) {
	w.ResponseWriter.Header().Del("Content-Length")
	w.ResponseWriter.WriteHeader(code)
}

// Gzip compresses all responses when the client sends Accept-Encoding: gzip.
// Uses a pool of gzip writers to avoid per-request allocation overhead.
func Gzip() gin.HandlerFunc {
	return func(c *gin.Context) {
		if !strings.Contains(c.GetHeader("Accept-Encoding"), "gzip") {
			c.Next()
			return
		}

		orig := c.Writer
		gz := gzPool.Get().(*gzip.Writer)
		gz.Reset(orig) // compressed output goes into the original gin writer

		c.Header("Content-Encoding", "gzip")
		c.Header("Vary", "Accept-Encoding")
		orig.Header().Del("Content-Length")

		c.Writer = &gzResponseWriter{orig, gz}

		defer func() {
			gz.Close()
			gzPool.Put(gz)
		}()

		c.Next()
	}
}