types.go

package compression

import (
	"io"
	"net/http"
	"sync"
)

// CompressionType represents the type of compression to use.
//
// The framework supports multiple compression algorithms, each with different
// characteristics in terms of compression ratio, speed, and browser support.
//
// Example:
//
//	// Use gzip compression (widely supported)
//	app.Use(compression.Compress(compression.UnifiedConfig{
//		Types: []compression.CompressionType{compression.GzipCompression},
//	}))
//
//	// Use brotli compression (better compression, modern browsers)
//	app.Use(compression.Compress(compression.UnifiedConfig{
//		Types: []compression.CompressionType{compression.BrotliCompression},
//	}))
//
//	// Prefer brotli, fallback to gzip
//	app.Use(compression.Compress(compression.UnifiedConfig{
//		Types: []compression.CompressionType{
//			compression.BrotliCompression,
//			compression.GzipCompression,
//		},
//	}))
type CompressionType string

const (
	// GzipCompression uses gzip compression algorithm.
	//
	// Gzip is the most widely supported compression format, available in all
	// modern browsers and HTTP clients. It provides good compression ratios
	// with reasonable CPU usage.
	//
	// Compression levels: 1-9 (higher = better compression but slower)
	// Browser support: Universal (all modern browsers)
	// Compression ratio: Good (typically 60-80% for text content)
	//
	// Example:
	//
	//	app.Use(compression.Gzip(compression.CompressionConfig{
	//		Level: 6, // balanced performance
	//	}))
	GzipCompression CompressionType = "gzip"

	// BrotliCompression uses brotli compression algorithm.
	//
	// Brotli is a modern compression algorithm that typically provides better
	// compression ratios than gzip, especially for text content like HTML,
	// CSS, and JavaScript. It's supported by all modern browsers.
	//
	// Compression levels: 0-11 (higher = better compression but slower)
	// Browser support: Modern browsers (Chrome 50+, Firefox 44+, Safari 11+)
	// Compression ratio: Excellent (typically 70-90% for text content)
	//
	// Note: Requires the github.com/andybalholm/brotli package to be available.
	// If not available, the middleware will gracefully fall back to no compression.
	//
	// Example:
	//
	//	app.Use(compression.Brotli(compression.CompressionConfig{
	//		Level: 11, // maximum compression
	//	}))
	BrotliCompression CompressionType = "br"
)

// CompressionConfig provides common configuration for all compression types.
//
// This configuration is used by individual compression middlewares (Gzip, Brotli)
// to control compression behavior. For unified compression with multiple types,
// use UnifiedConfig instead.
//
// Example:
//
//	// Basic gzip compression
//	app.Use(compression.Gzip(compression.CompressionConfig{
//		Level: 6,
//		MinSize: 1024,
//	}))
//
//	// High compression for static assets
//	app.GET("/static/*", staticHandler, compression.Gzip(compression.CompressionConfig{
//		Level: 9,    // maximum compression
//		MinSize: 256, // compress even small files
//	}))
type CompressionConfig struct {
	// Level sets the compression level.
	//
	// For gzip: 1-9 (higher = better compression but slower)
	//   - 1: Fastest compression, lowest ratio
	//   - 6: Balanced performance (default)
	//   - 9: Best compression, slowest
	//
	// For brotli: 0-11 (higher = better compression but slower)
	//   - 0: Fastest compression, lowest ratio
	//   - 6: Balanced performance (default)
	//   - 11: Best compression, slowest
	//
	// Default: 6 (balanced performance)
	Level int

	// MinSize sets the minimum response size to compress (in bytes).
	//
	// Responses smaller than this threshold are not compressed to avoid
	// the overhead of compression setup and headers. This is especially
	// important for small API responses where compression overhead
	// might exceed the benefits.
	//
	// Default: 0 (compress all responses)
	// Recommended: 1024 bytes for general use, 256 bytes for text content
	MinSize int
}

// CompressionWriter represents a compression writer with cleanup functionality.
//
// This struct wraps an io.WriteCloser with a cleanup function that should
// be called when the writer is no longer needed. The cleanup function
// typically returns the writer to a pool for reuse.
//
// Example:
//
//	writer, err := provider.GetWriter(level, responseWriter)
//	if err != nil {
//		return err
//	}
//	defer writer.Put() // Return to pool
//
//	// Use writer.Writer for compression
//	_, err = writer.Writer.Write(data)
type CompressionWriter struct {
	// Writer is the underlying compression writer.
	// Use this for writing compressed data to the response.
	Writer io.WriteCloser

	// Put is a function that should be called to return the writer to the pool.
	// This ensures proper cleanup and prevents memory leaks.
	Put func()
}

// CompressionProvider defines the interface for compression implementations.
//
// This interface allows the compression package to work with different
// compression algorithms in a unified way. Each compression type (gzip,
// brotli, etc.) implements this interface.
//
// The interface provides methods for:
//   - Creating new compression writers
//   - Getting pooled compression writers (for performance)
//   - Checking if the compression type is supported
//   - Getting the compression type identifier
//
// Example:
//
//	// Create a gzip provider
//	gzipProvider := compression.NewGzipProvider()
//
//	// Check if supported (gzip is always supported)
//	if gzipProvider.IsSupported() {
//		// Use the provider
//		writer, err := gzipProvider.GetWriter(6, responseWriter)
//		if err != nil {
//			return err
//		}
//		defer writer.Put()
//	}
type CompressionProvider interface {
	// NewWriter creates a new compression writer for the given output and level.
	//
	// This method creates a fresh compression writer without pooling.
	// For better performance, prefer GetWriter which uses pooled writers.
	//
	// Args:
	//   - w: The output writer to compress to
	//   - level: The compression level (1-9 for gzip, 0-11 for brotli)
	//
	// Returns:
	//   - io.WriteCloser: The compression writer
	//   - error: Any error that occurred during creation
	NewWriter(w io.Writer, level int) (io.WriteCloser, error)

	// GetWriter returns a pooled compression writer for the given level and output.
	//
	// This method reuses compression writers from a pool to minimize
	// allocations and improve performance. The returned CompressionWriter
	// includes a Put function that should be called to return the writer
	// to the pool.
	//
	// Args:
	//   - level: The compression level (1-9 for gzip, 0-11 for brotli)
	//   - w: The output writer to compress to
	//
	// Returns:
	//   - *CompressionWriter: Pooled compression writer with cleanup function
	//   - error: Any error that occurred during creation
	GetWriter(level int, w io.Writer) (*CompressionWriter, error)

	// CompressionType returns the type of compression this provider implements.
	//
	// This method returns the CompressionType constant that identifies
	// the compression algorithm (e.g., GzipCompression, BrotliCompression).
	CompressionType() CompressionType

	// IsSupported returns whether this compression type is supported.
	//
	// This method checks if the compression algorithm is available
	// on the current system. Gzip is always supported (stdlib),
	// while brotli requires the external package to be available.
	//
	// Returns:
	//   - bool: true if the compression type is supported, false otherwise
	IsSupported() bool
}

// ResponseWriter wraps http.ResponseWriter to provide compression.
//
// This struct implements http.ResponseWriter and adds compression capabilities.
// It automatically handles:
//   - Compression decision based on Accept-Encoding header
//   - Compression level and minimum size thresholds
//   - Proper HTTP headers (Content-Encoding, Vary)
//   - Writer pooling for performance
//   - Automatic cleanup via defer
//
// The compression is applied lazily - the compression writer is only
// created when the first Write call is made, avoiding unnecessary
// setup for responses that might be empty.
//
// Example:
//
//	// In middleware
//	crw := &compression.ResponseWriter{
//		rw:       c.ResponseWriter(),
//		provider: compression.NewGzipProvider(),
//		level:    6,
//		minSize:  1024,
//	}
//	c.SetResponseWriter(crw)
//	defer crw.Close() // Ensure cleanup
//
//	// Continue with handler - compression happens automatically
//	return next(c)
type ResponseWriter struct {
	rw             http.ResponseWriter
	provider       CompressionProvider
	level          int
	minSize        int
	writer         io.WriteCloser
	put            func()
	wroteHeader    bool
	useCompression bool
	contentLength  int
}

// Header returns the underlying response headers map.
//
// This method provides access to the response headers for setting
// custom headers before compression is applied.
//
// Example:
//
//	crw := &compression.ResponseWriter{...}
//	crw.Header().Set("Content-Type", "application/json")
//	crw.Header().Set("Cache-Control", "public, max-age=3600")
func (r *ResponseWriter) Header() http.Header { return r.rw.Header() }

// WriteHeader writes the HTTP response header and decides whether to use compression.
//
// This method is called automatically when the first Write call is made,
// or can be called explicitly to control when headers are written.
//
// Compression decision logic:
//   - Skips compression for HEAD requests (no body to compress)
//   - Skips compression for 204 No Content and 304 Not Modified responses
//   - Skips compression for 1xx informational responses
//   - Skips compression if Content-Encoding is already set
//   - Skips compression if the provider is not supported
//   - Sets appropriate headers (Content-Encoding, Vary) if compressing
//
// Example:
//
//	crw := &compression.ResponseWriter{...}
//	crw.Header().Set("Content-Type", "application/json")
//	crw.WriteHeader(http.StatusOK) // Headers written, compression decided
func (r *ResponseWriter) WriteHeader(status int) {
	if r.wroteHeader {
		return
	}
	r.wroteHeader = true

	// Skip compression for informational responses (1xx)
	if status < 200 {
		r.useCompression = false
		r.rw.WriteHeader(status)
		return
	}

	// Skip compression for certain status codes that shouldn't have a body
	if status == http.StatusNoContent || status == http.StatusNotModified {
		r.useCompression = false
		r.rw.WriteHeader(status)
		return
	}

	// Check if content is already encoded
	enc := r.Header().Get("Content-Encoding")
	if enc != "" && enc != "identity" {
		r.useCompression = false
		r.rw.WriteHeader(status)
		return
	}

	// Only set compression headers if the provider is supported
	if r.provider != nil && r.provider.IsSupported() {
		r.useCompression = true
		r.Header().Del("Content-Length")
		r.Header().Set("Content-Encoding", string(r.provider.CompressionType()))
		r.Header().Add("Vary", "Accept-Encoding")
	} else {
		r.useCompression = false
	}
	r.rw.WriteHeader(status)
}

// Write writes data to the response, compressing if enabled.
//
// This method handles the compression logic:
//   - Automatically calls WriteHeader if not already called
//   - Tracks content length for minimum size threshold
//   - Disables compression if response becomes too small
//   - Creates compression writer lazily on first write
//   - Handles compression errors gracefully
//
// The compression writer is created lazily to avoid unnecessary
// setup for responses that might be empty or very small.
//
// Example:
//
//	crw := &compression.ResponseWriter{...}
//	data := []byte("Hello, World!")
//	_, err := crw.Write(data) // Compression happens automatically
//	if err != nil {
//		return err
//	}
func (r *ResponseWriter) Write(p []byte) (int, error) {
	if !r.wroteHeader {
		r.WriteHeader(http.StatusOK)
	}

	r.contentLength += len(p)

	// Don't compress if response is too small
	if r.useCompression && r.contentLength < r.minSize {
		r.useCompression = false
		// Remove compression headers if we're not compressing
		r.Header().Del("Content-Encoding")
		r.Header().Del("Vary")
	}

	if !r.useCompression {
		return r.rw.Write(p)
	}

	if r.writer == nil {
		compWriter, err := r.provider.GetWriter(r.level, r.rw)
		if err != nil {
			// Fallback to no compression on error
			r.useCompression = false
			r.Header().Del("Content-Encoding")
			r.Header().Del("Vary")
			return r.rw.Write(p)
		}
		r.writer = compWriter.Writer
		r.put = compWriter.Put
	}

	return r.writer.Write(p)
}

// Close ensures proper cleanup of the compression writer.
//
// This method should be called via defer in middleware to ensure
// proper cleanup and prevent memory leaks. It handles:
//   - Returning the writer to the pool (if pooled)
//   - Closing the compression writer
//   - Cleaning up internal state
//
// Example:
//
//	// In middleware
//	crw := &compression.ResponseWriter{...}
//	c.SetResponseWriter(crw)
//	defer crw.Close() // Always call this
//
//	return next(c)
func (r *ResponseWriter) Close() error {
	if r.writer != nil {
		if r.put != nil {
			r.put()
			r.writer, r.put = nil, nil
			return nil
		}
		return r.writer.Close()
	}
	return nil
}

// Flush flushes the compression writer and forwards to the underlying ResponseWriter.
//
// This method is useful for streaming responses or when immediate
// data transmission is needed. It flushes both the compression
// writer (if it supports flushing) and the underlying ResponseWriter.
//
// Example:
//
//	crw := &compression.ResponseWriter{...}
//	crw.Write([]byte("partial data"))
//	crw.Flush() // Ensure data is sent immediately
//	crw.Write([]byte("more data"))
func (r *ResponseWriter) Flush() {
	if r.writer != nil {
		if f, ok := r.writer.(interface{ Flush() error }); ok {
			_ = f.Flush()
		}
	}
	if f, ok := r.rw.(http.Flusher); ok {
		f.Flush()
	}
}

// Ensure interfaces are implemented
var _ http.ResponseWriter = (*ResponseWriter)(nil)
var _ http.Flusher = (*ResponseWriter)(nil)

// Global compression pools for reuse across different compression types
var compressionPools sync.Map // map[string]*sync.Pool

// getPool returns a sync.Pool for the given compression type and level.
//
// This function manages global compression writer pools to minimize
// allocations and improve performance. Each compression type and level
// combination gets its own pool.
//
// Args:
//   - compType: The compression type (gzip, brotli, etc.)
//   - level: The compression level
//
// Returns:
//   - *sync.Pool: The pool for the given type and level
//
// Example:
//
//	// Get a pool for gzip level 6
//	pool := getPool(compression.GzipCompression, 6)
//
//	// Get a pool for brotli level 11
//	pool := getPool(compression.BrotliCompression, 11)
func getPool(compType CompressionType, level int) *sync.Pool {
	key := string(compType) + ":" + string(rune(level))
	poolAny, _ := compressionPools.LoadOrStore(key, &sync.Pool{})
	return poolAny.(*sync.Pool)
}