package server

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

type GZipResponseWriter struct {
	w  http.ResponseWriter
	gz *gzip.Writer
}

func NewGZipResponseWriter(w http.ResponseWriter) *GZipResponseWriter {
	w.Header().Add("Content-Type", "text/html")
	w.Header().Add("Content-Encoding", "gzip")
	return &GZipResponseWriter{
		w:  w,
		gz: gzip.NewWriter(w),
	}
}

func (gz *GZipResponseWriter) Header() http.Header {
	return gz.w.Header()
}

func (gz *GZipResponseWriter) WriteHeader(statusCode int) {
	gz.w.WriteHeader(statusCode)
}

func (gz *GZipResponseWriter) Write(b []byte) (int, error) {
	return gz.gz.Write(b)
}

func (gz *GZipResponseWriter) Close() error {
	return gz.gz.Close()
}

func (s *Server) gzip(h http.HandlerFunc) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		if encoding, _ := r.Header["Accept-Encoding"]; strings.Contains(fmt.Sprint(encoding), "gzip") {
			gz := NewGZipResponseWriter(w)
			defer gz.Close()
			w = gz
		}
		h(w, r)
	}
}