package handlers

import (
	"compress/gzip"
	"net/http"
	"strings"
)

// Gzip 对 JSON/文本响应启用 gzip 压缩（仅客户端声明 Accept-Encoding: gzip 时）。
// 图片/音频等二进制响应不在白名单内，原样透传。
func Gzip(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if !strings.Contains(r.Header.Get("Accept-Encoding"), "gzip") {
			next.ServeHTTP(w, r)
			return
		}
		gw := &gzipWriter{ResponseWriter: w}
		defer gw.close()
		next.ServeHTTP(gw, r)
	})
}

type gzipWriter struct {
	http.ResponseWriter
	zw         *gzip.Writer
	decided    bool // 是否已根据 Content-Type 做过压缩决定
	headerSent bool
}

func (g *gzipWriter) WriteHeader(status int) {
	if g.headerSent {
		return
	}
	if !g.decided {
		g.decide()
	}
	g.ResponseWriter.WriteHeader(status)
	g.headerSent = true
}

func (g *gzipWriter) Write(b []byte) (int, error) {
	if !g.headerSent {
		// handler 未显式 WriteHeader（fmt.Fprint 直接写）时的隐式 200
		if !g.decided {
			g.decide()
		}
		g.ResponseWriter.WriteHeader(http.StatusOK)
		g.headerSent = true
	}
	if g.zw != nil {
		return g.zw.Write(b)
	}
	return g.ResponseWriter.Write(b)
}

// decide 在首个响应头确定时判断是否压缩（此时 handler 已设置 Content-Type）
func (g *gzipWriter) decide() {
	g.decided = true
	ct := g.Header().Get("Content-Type")
	for _, prefix := range []string{"application/json", "text/html", "text/plain"} {
		if strings.HasPrefix(ct, prefix) {
			g.Header().Set("Content-Encoding", "gzip")
			g.Header().Del("Content-Length")
			g.zw = gzip.NewWriter(g.ResponseWriter)
			return
		}
	}
}

func (g *gzipWriter) close() {
	if g.zw != nil {
		g.zw.Close()
	}
}