package middleware

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

type gzipResponseWriter struct {
	http.ResponseWriter
	writer io.Writer
	wrote  bool
}

func (w *gzipResponseWriter) WriteHeader(status int) {
	if w.wrote {
		return
	}
	w.wrote = true
	header := w.Header()
	header.Del("Content-Length")
	header.Del("Accept-Ranges")
	header.Set("Content-Encoding", "gzip")
	header.Add("Vary", "Accept-Encoding")
	w.ResponseWriter.WriteHeader(status)
}

func (w *gzipResponseWriter) Write(data []byte) (int, error) {
	if !w.wrote {
		w.WriteHeader(http.StatusOK)
	}
	return w.writer.Write(data)
}

func Gzip(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if !acceptsGzip(r) || shouldSkipGzip(r) {
			next.ServeHTTP(w, r)
			return
		}

		gz := gzip.NewWriter(w)
		defer gz.Close()

		next.ServeHTTP(&gzipResponseWriter{
			ResponseWriter: w,
			writer:         gz,
		}, r)
	})
}

func acceptsGzip(r *http.Request) bool {
	return strings.Contains(r.Header.Get("Accept-Encoding"), "gzip")
}

func shouldSkipGzip(r *http.Request) bool {
	if r.Method == http.MethodHead || r.Header.Get("Range") != "" {
		return true
	}

	path := r.URL.Path
	switch ext := strings.ToLower(path[strings.LastIndex(path, ".")+1:]); ext {
	case "br", "gz", "zip", "7z", "rar", "jpg", "jpeg", "png", "gif", "webp", "avif", "mp4", "webm", "mov", "m4v", "mp3", "ogg", "woff", "woff2", "ttf", "otf":
		return true
	default:
		return false
	}
}