package router

import (
	"bufio"
	"bytes"
	"compress/gzip"
	"io"
	"net"
	"net/http"
	"strconv"
	"strings"

	"github.com/gin-gonic/gin"
)

const gzipCompressionMinBytes = 1024

func gzipCompressionMiddleware() gin.HandlerFunc {
	return func(ctx *gin.Context) {
		if !requestAcceptsGzip(ctx.Request) || isWebSocketUpgrade(ctx.Request) {
			ctx.Next()
			return
		}

		original := ctx.Writer
		writer := newGzipResponseWriter(ctx.Writer)
		ctx.Writer = writer
		defer func() {
			ctx.Writer = original
		}()
		ctx.Next()
		writer.finish()
	}
}

type gzipResponseWriter struct {
	gin.ResponseWriter
	status      int
	size        int
	wroteHeader bool
	passthrough bool
	compressing bool
	gzipWriter  *gzip.Writer
	buffer      bytes.Buffer
}

func newGzipResponseWriter(writer gin.ResponseWriter) *gzipResponseWriter {
	return &gzipResponseWriter{ResponseWriter: writer, status: http.StatusOK}
}

func (w *gzipResponseWriter) WriteHeader(code int) {
	// Gin uses non-positive status codes to render without committing headers.
	if code <= 0 || w.wroteHeader {
		return
	}
	w.status = code
	w.wroteHeader = true
	if w.passthrough {
		w.ResponseWriter.WriteHeader(code)
	}
}

func (w *gzipResponseWriter) WriteHeaderNow() {
	if !w.wroteHeader {
		w.WriteHeader(w.status)
	}
	if w.passthrough {
		w.ResponseWriter.WriteHeaderNow()
	}
}

func (w *gzipResponseWriter) Write(data []byte) (int, error) {
	if !w.wroteHeader {
		w.WriteHeader(http.StatusOK)
	}
	if w.shouldBypassFromHeaders() {
		w.enablePassthrough()
	}
	if w.passthrough {
		return w.ResponseWriter.Write(data)
	}
	w.size += len(data)
	if w.compressing {
		if _, err := w.gzipWriter.Write(data); err != nil {
			return 0, err
		}
		return len(data), nil
	}
	_, _ = w.buffer.Write(data)
	if w.buffer.Len() < gzipCompressionMinBytes {
		return len(data), nil
	}
	if !w.shouldCompress() {
		w.enablePassthrough()
		return len(data), nil
	}
	if err := w.enableCompression(); err != nil {
		return 0, err
	}
	return len(data), nil
}

func (w *gzipResponseWriter) WriteString(data string) (int, error) {
	return w.Write([]byte(data))
}

func (w *gzipResponseWriter) Status() int {
	if w.wroteHeader {
		return w.status
	}
	return w.ResponseWriter.Status()
}

func (w *gzipResponseWriter) Size() int {
	return w.size
}

func (w *gzipResponseWriter) Written() bool {
	return w.wroteHeader || w.ResponseWriter.Written()
}

func (w *gzipResponseWriter) Flush() {
	if w.compressing {
		_ = w.gzipWriter.Flush()
		w.ResponseWriter.Flush()
		return
	}
	w.enablePassthrough()
	w.ResponseWriter.Flush()
}

func (w *gzipResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
	w.enablePassthrough()
	return w.ResponseWriter.Hijack()
}

func (w *gzipResponseWriter) CloseNotify() <-chan bool {
	return w.ResponseWriter.CloseNotify()
}

func (w *gzipResponseWriter) Pusher() http.Pusher {
	return w.ResponseWriter.Pusher()
}

func (w *gzipResponseWriter) finish() {
	if w.passthrough {
		return
	}
	if w.compressing {
		_ = w.gzipWriter.Close()
		return
	}
	if !w.wroteHeader && w.buffer.Len() == 0 {
		return
	}
	if !w.wroteHeader {
		w.WriteHeader(http.StatusOK)
	}
	if !w.shouldCompress() {
		w.ResponseWriter.WriteHeader(w.status)
		_, _ = w.ResponseWriter.Write(w.buffer.Bytes())
		return
	}

	if err := w.enableCompression(); err != nil {
		return
	}
	_ = w.gzipWriter.Close()
}

func (w *gzipResponseWriter) enableCompression() error {
	if w.compressing {
		return nil
	}
	gzipWriter, err := gzip.NewWriterLevel(w.ResponseWriter, gzip.DefaultCompression)
	if err != nil {
		return err
	}
	header := w.Header()
	appendVary(header, "Accept-Encoding")
	header.Del("Content-Length")
	header.Set("Content-Encoding", "gzip")
	w.ResponseWriter.WriteHeader(w.status)
	w.gzipWriter = gzipWriter
	w.compressing = true
	if w.buffer.Len() == 0 {
		return nil
	}
	_, err = io.Copy(gzipWriter, bytes.NewReader(w.buffer.Bytes()))
	w.buffer.Reset()
	return err
}

func (w *gzipResponseWriter) enablePassthrough() {
	if w.passthrough {
		return
	}
	if w.compressing {
		return
	}
	w.passthrough = true
	if w.wroteHeader {
		w.ResponseWriter.WriteHeader(w.status)
	}
	if w.buffer.Len() > 0 {
		_, _ = w.ResponseWriter.Write(w.buffer.Bytes())
		w.buffer.Reset()
	}
}

func (w *gzipResponseWriter) shouldBypassFromHeaders() bool {
	header := w.Header()
	if strings.TrimSpace(header.Get("Content-Encoding")) != "" {
		return true
	}
	if hasCacheControlDirective(header, "no-transform") {
		return true
	}
	if strings.EqualFold(strings.TrimSpace(header.Get("X-Accel-Buffering")), "no") {
		return true
	}
	contentType := normalizeContentType(header.Get("Content-Type"))
	return isCompressionExcludedContentType(contentType)
}

func (w *gzipResponseWriter) shouldCompress() bool {
	if w.buffer.Len() < gzipCompressionMinBytes {
		return false
	}
	if w.status < 200 || w.status == http.StatusNoContent || w.status == http.StatusNotModified {
		return false
	}
	header := w.Header()
	if strings.TrimSpace(header.Get("Content-Encoding")) != "" {
		return false
	}
	if hasCacheControlDirective(header, "no-transform") {
		return false
	}
	if strings.EqualFold(strings.TrimSpace(header.Get("X-Accel-Buffering")), "no") {
		return false
	}
	return isCompressibleContentType(normalizeContentType(header.Get("Content-Type")))
}

func requestAcceptsGzip(request *http.Request) bool {
	if request == nil || request.Method == http.MethodHead {
		return false
	}
	for _, value := range strings.Split(request.Header.Get("Accept-Encoding"), ",") {
		parts := strings.Split(strings.TrimSpace(value), ";")
		if !strings.EqualFold(strings.TrimSpace(parts[0]), "gzip") {
			continue
		}
		for _, param := range parts[1:] {
			keyValue := strings.SplitN(strings.TrimSpace(param), "=", 2)
			if len(keyValue) != 2 || !strings.EqualFold(strings.TrimSpace(keyValue[0]), "q") {
				continue
			}
			q, err := strconv.ParseFloat(strings.TrimSpace(keyValue[1]), 64)
			if err == nil && q <= 0 {
				return false
			}
		}
		return true
	}
	return false
}

func isWebSocketUpgrade(request *http.Request) bool {
	if request == nil {
		return false
	}
	if !strings.EqualFold(strings.TrimSpace(request.Header.Get("Upgrade")), "websocket") {
		return false
	}
	for _, part := range strings.Split(request.Header.Get("Connection"), ",") {
		if strings.EqualFold(strings.TrimSpace(part), "upgrade") {
			return true
		}
	}
	return false
}

func normalizeContentType(contentType string) string {
	contentType = strings.ToLower(strings.TrimSpace(contentType))
	if index := strings.Index(contentType, ";"); index >= 0 {
		contentType = strings.TrimSpace(contentType[:index])
	}
	return contentType
}

func isCompressionExcludedContentType(contentType string) bool {
	switch contentType {
	case "text/event-stream", "application/x-ndjson":
		return true
	default:
		return false
	}
}

func isCompressibleContentType(contentType string) bool {
	if contentType == "" || isCompressionExcludedContentType(contentType) {
		return false
	}
	if strings.HasPrefix(contentType, "text/") {
		return true
	}
	switch contentType {
	case "application/json", "application/javascript", "application/x-javascript", "application/xml", "image/svg+xml":
		return true
	default:
		return strings.HasSuffix(contentType, "+json") || strings.HasSuffix(contentType, "+xml")
	}
}

func hasCacheControlDirective(header http.Header, directive string) bool {
	for _, value := range header.Values("Cache-Control") {
		for _, part := range strings.Split(value, ",") {
			name := strings.TrimSpace(strings.SplitN(part, "=", 2)[0])
			if strings.EqualFold(name, directive) {
				return true
			}
		}
	}
	return false
}

func appendVary(header http.Header, value string) {
	for _, existing := range header.Values("Vary") {
		for _, part := range strings.Split(existing, ",") {
			if strings.EqualFold(strings.TrimSpace(part), value) {
				return
			}
		}
	}
	header.Add("Vary", value)
}