// Copyright 2026 tamakyi. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.

package route

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

// GzipHandler 对可压缩的文本响应（HTML / CSS / JS / JSON 等）做 gzip 压缩，
// 自动跳过图片、视频、字体等已压缩格式。
// flamego 的 Context 无法替换底层 ResponseWriter，故在 http.Server 最外层包装
// Handler，保证对所有响应（含静态文件）生效。
func GzipHandler(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 := &gzipResponseWriter{
			ResponseWriter: w,
			gz:             gzip.NewWriter(w),
		}
		defer func() {
			// 只有真正写入过 body 才关闭，避免 302/204 等无 body 响应被追加空的 gzip 尾块
			if gw.compress && gw.wrote {
				_ = gw.gz.Close()
			}
		}()
		next.ServeHTTP(gw, r)
	})
}

// gzipResponseWriter 在 WriteHeader 时按 Content-Type 决定是否启用 gzip，
// 并保持 Flusher / Hijacker / Pusher 接口兼容。
type gzipResponseWriter struct {
	http.ResponseWriter
	gz          *gzip.Writer
	compress    bool
	wrote       bool
	wroteHeader bool
}

// WriteHeader 记录状态码并决定是否压缩；Content-Encoding 改变后 Content-Length
// 不再准确，需删除该头。
func (w *gzipResponseWriter) WriteHeader(status int) {
	if w.wroteHeader {
		return
	}
	w.wroteHeader = true

	// 无实体响应（204 No Content / 304 Not Modified）不应附加 Content-Encoding。
	if status == http.StatusNoContent || status == http.StatusNotModified {
		w.ResponseWriter.WriteHeader(status)
		return
	}

	if shouldCompressContentType(w.Header().Get("Content-Type")) {
		w.compress = true
		w.Header().Del("Content-Length")
		w.Header().Set("Content-Encoding", "gzip")
		w.Header().Add("Vary", "Accept-Encoding")
	}
	w.ResponseWriter.WriteHeader(status)
}

// Write 写入响应体。若尚未 WriteHeader，先用首个字节块探测类型再决策。
func (w *gzipResponseWriter) Write(b []byte) (int, error) {
	if !w.wroteHeader {
		ct := w.Header().Get("Content-Type")
		if ct == "" {
			ct = http.DetectContentType(b)
			// 嗅探结果须写回头，否则空 Content-Type 会被判为「不压缩」，
			// 导致本应压缩的纯文本响应漏压。
			w.Header().Set("Content-Type", ct)
		}
		w.WriteHeader(http.StatusOK)
	}
	if w.compress {
		w.wrote = true
		return w.gz.Write(b)
	}
	return w.ResponseWriter.Write(b)
}

// Flush 转发底层 Flusher，压缩模式下先刷新 gzip 缓冲。
func (w *gzipResponseWriter) Flush() {
	if w.compress && w.gz != nil {
		_ = w.gz.Flush()
	}
	if f, ok := w.ResponseWriter.(http.Flusher); ok {
		f.Flush()
	}
}

// Hijack 转发底层 Hijacker，保持对 WebSocket 等升级协议兼容。
func (w *gzipResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
	if h, ok := w.ResponseWriter.(http.Hijacker); ok {
		return h.Hijack()
	}
	return nil, nil, errors.New("gzipResponseWriter: hijacking not supported")
}

// Push 转发 HTTP/2 Server Push（如果底层支持）。
func (w *gzipResponseWriter) Push(target string, opts *http.PushOptions) error {
	if p, ok := w.ResponseWriter.(http.Pusher); ok {
		return p.Push(target, opts)
	}
	return http.ErrNotSupported
}

// shouldCompressContentType 判断 Content-Type 是否值得压缩：空类型与已压缩的
// 二进制格式（图片 / 视频 / 音频 / 字体 / 归档等）直接跳过。
func shouldCompressContentType(ct string) bool {
	ct = strings.ToLower(strings.TrimSpace(ct))
	if ct == "" {
		return false
	}

	if strings.HasPrefix(ct, "image/") {
		return strings.Contains(ct, "svg") // SVG 是文本格式
	}
	if strings.HasPrefix(ct, "video/") || strings.HasPrefix(ct, "audio/") ||
		strings.HasPrefix(ct, "font/") {
		return false
	}
	switch ct {
	case "application/zip", "application/gzip", "application/x-7z-compressed",
		"application/x-rar-compressed", "application/pdf",
		"application/vnd.ms-fontobject", "application/wasm",
		"application/octet-stream", "application/x-protobuf", "application/grpc":
		return false
	}
	return true
}