package server

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

const gzipMin = 8 << 10

type gzipWriter struct {
	http.ResponseWriter
	held int
	gz   *gzip.Writer
}

func (w *gzipWriter) WriteHeader(status int) {
	if status == http.StatusOK && w.Header().Get("Content-Type") == "application/json" {
		w.held = status
		return
	}
	w.ResponseWriter.WriteHeader(status)
}

func (w *gzipWriter) Write(b []byte) (int, error) {
	if w.held != 0 {
		if len(b) >= gzipMin {
			w.Header().Set("Content-Encoding", "gzip")
			w.Header().Del("Content-Length")
			w.gz = gzip.NewWriter(w.ResponseWriter)
		}
		w.ResponseWriter.WriteHeader(w.held)
		w.held = 0
	}
	if w.gz != nil {
		return w.gz.Write(b)
	}
	return w.ResponseWriter.Write(b)
}

func (w *gzipWriter) Unwrap() http.ResponseWriter {
	return w.ResponseWriter
}

func compress(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}
		next.ServeHTTP(gw, r)
		if gw.held != 0 {
			w.WriteHeader(gw.held)
		}
		if gw.gz != nil {
			err := gw.gz.Close()
			if err != nil {
				slog.Warn("write response", "err", err)
			}
		}
	})
}