package middleware

import (
	"compress/flate"
	"compress/gzip"
	"strings"

	"github.com/andybalholm/brotli"
	"github.com/gogf/gf/v2/net/ghttp"
)

func CompressMiddleware(r *ghttp.Request) {
	ae := strings.ToLower(r.Header.Get("Accept-Encoding"))

	// 先执行后面的逻辑
	r.Middleware.Next()

	// 已经压缩过就跳过
	if r.Response.Header().Get("Content-Encoding") != "" {
		return
	}

	// 只压缩文本类型
	contentType := r.Response.Header().Get("Content-Type")
	compressTypes := []string{
		"text/",
		"application/json",
		"application/javascript",
		"application/xml",
	}
	need := false
	for _, t := range compressTypes {
		if strings.HasPrefix(contentType, t) {
			need = true
			break
		}
	}
	if !need {
		return
	}

	buf := r.Response.Buffer()
	if len(buf) == 0 {
		return
	}

	// 优先 Brotli
	if strings.Contains(ae, "br") {
		r.Response.Header().Set("Content-Encoding", "br")
		r.Response.Header().Del("Content-Length")
		r.Response.ClearBuffer()

		wr := brotli.NewWriterLevel(r.Response.Writer, brotli.BestCompression)
		defer wr.Close()
		wr.Write(buf)
		return
	}

	// 降级 Gzip
	if strings.Contains(ae, "gzip") {
		r.Response.Header().Set("Content-Encoding", "gzip")
		r.Response.Header().Del("Content-Length")
		r.Response.ClearBuffer()

		wr, _ := gzip.NewWriterLevel(r.Response.Writer, flate.BestCompression)
		defer wr.Close()
		wr.Write(buf)
		return
	}
}