package middleware

import (
	"compress/gzip"
	"io"
	"net/http"
	"strings"
	"sync"

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

// gzipMinLength 小于此字节数的响应不压缩（避免小响应压缩反而增大体积）
const gzipMinLength = 1024

// gzipPool 复用 gzip.Writer，减少 GC 压力
var gzipPool = sync.Pool{
	New: func() interface{} {
		w, _ := gzip.NewWriterLevel(io.Discard, gzip.BestSpeed)
		return w
	},
}

// gzipResponseWriter 包装 gin.ResponseWriter，按需启用 gzip 压缩。
// 所有写入先缓冲；第一次超过阈值时决定是否压缩。
type gzipResponseWriter struct {
	gin.ResponseWriter
	gz        *gzip.Writer
	buf       []byte
	threshold int
	decided   bool // 已决定是否压缩
}

func (g *gzipResponseWriter) Write(data []byte) (int, error) {
	if g.decided {
		if g.gz != nil {
			return g.gz.Write(data)
		}
		return g.ResponseWriter.Write(data)
	}

	// 积累数据
	g.buf = append(g.buf, data...)
	if len(g.buf) >= g.threshold {
		return len(data), g.decide()
	}
	return len(data), nil
}

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

// decide 根据已缓冲内容和 Content-Type 决定是否压缩，并写出缓冲数据
func (g *gzipResponseWriter) decide() error {
	g.decided = true

	ct := g.ResponseWriter.Header().Get("Content-Type")
	if g.gz != nil && shouldCompress(ct) {
		// 启用 gzip
		g.ResponseWriter.Header().Set("Content-Encoding", "gzip")
		g.ResponseWriter.Header().Set("Vary", "Accept-Encoding")
		g.ResponseWriter.Header().Del("Content-Length")
		g.gz.Reset(g.ResponseWriter)
		if len(g.buf) > 0 {
			_, err := g.gz.Write(g.buf)
			g.buf = nil
			return err
		}
	} else {
		// 不压缩：回收 gzip.Writer
		if g.gz != nil {
			gzipPool.Put(g.gz)
			g.gz = nil
		}
		if len(g.buf) > 0 {
			_, err := g.ResponseWriter.Write(g.buf)
			g.buf = nil
			return err
		}
	}
	g.buf = nil
	return nil
}

// finalize 在请求处理完毕后刷出剩余缓冲数据并关闭 gzip.Writer
func (g *gzipResponseWriter) finalize() {
	if !g.decided {
		// 响应体小于阈值，直接透传（不压缩）
		g.decided = true
		if g.gz != nil {
			gzipPool.Put(g.gz)
			g.gz = nil
		}
		if len(g.buf) > 0 {
			_, _ = g.ResponseWriter.Write(g.buf)
			g.buf = nil
		}
		return
	}

	if g.gz != nil {
		_ = g.gz.Flush()
		_ = g.gz.Close()
		gzipPool.Put(g.gz)
		g.gz = nil
	}
}

// shouldCompress 根据 Content-Type 判断是否值得压缩（二进制流不压缩）
func shouldCompress(contentType string) bool {
	ct := strings.ToLower(strings.SplitN(contentType, ";", 2)[0])
	switch ct {
	case "application/json",
		"application/javascript",
		"text/html",
		"text/plain",
		"text/css",
		"text/xml",
		"application/xml",
		"application/x-www-form-urlencoded":
		return true
	}
	return false
}

// GzipMiddleware 对 JSON/文本类响应启用 GZIP 压缩。
//
// 仅在满足以下条件时压缩：
//   - 客户端发送了 Accept-Encoding: gzip
//   - 响应 Content-Type 为 JSON/文本类
//   - 响应体超过 gzipMinLength（默认 1 KiB）
//
// 其余情况透传，不影响性能。
func GzipMiddleware() gin.HandlerFunc {
	return func(c *gin.Context) {
		// 客户端不接受 gzip 则跳过
		if !strings.Contains(c.GetHeader("Accept-Encoding"), "gzip") {
			c.Next()
			return
		}

		gz := gzipPool.Get().(*gzip.Writer)

		grw := &gzipResponseWriter{
			ResponseWriter: c.Writer,
			gz:             gz,
			threshold:      gzipMinLength,
		}

		c.Writer = grw

		defer func() {
			grw.finalize()
			c.Writer = grw.ResponseWriter
		}()

		c.Next()
	}
}

// Ensure gzipResponseWriter implements http.Hijacker forwarding (needed by some WebSocket libs)
var _ http.ResponseWriter = (*gzipResponseWriter)(nil)