package middleware

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

type gzipResponseWriter struct {
	gzip.Writer
	http.ResponseWriter
}

func newGzipResponseWriter(w http.ResponseWriter) *gzipResponseWriter {
	gzrw := new(gzipResponseWriter)
	gzrw.Writer = *gzip.NewWriter(w)
	gzrw.ResponseWriter = w

	return gzrw
}

func (gzrw *gzipResponseWriter) Close() error {
	return gzrw.Writer.Close()
}

func (gzrw *gzipResponseWriter) Header() http.Header {
	return gzrw.ResponseWriter.Header()
}

func (gzrw *gzipResponseWriter) Write(p []byte) (int, error) {
	n, err := gzrw.Writer.Write(p)
	return n, err
}

func (gzrw *gzipResponseWriter) WriteHeader(statuscode int) {
	gzrw.ResponseWriter.WriteHeader(statuscode)
}

// Middleware to conditionally gzip a response
func Gzip(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		acceptedEncodings := r.Header.Get("Accept-Encoding")

		if !strings.Contains(acceptedEncodings, "gzip") {
			Logger.DebugContext(r.Context(), "Using gzip middleware but client does not support gzip encoding")
			next.ServeHTTP(w, r)
			return
		}

		w.Header().Add("Content-Encoding", "gzip")
		gzrw := newGzipResponseWriter(w)
		defer gzrw.Close()

		next.ServeHTTP(gzrw, r)
	})
}