compression.go

package middleware

import (
	"bytes"
	"compress/gzip"
	"encoding/binary"
	"fmt"
	"io"
	"sync"

	"github.com/golang/snappy"
	"github.com/pierrec/lz4/v4"
)

type CompressionType int

const (
	CompressionNone CompressionType = iota
	CompressionGzip
	CompressionSnappy
	CompressionLZ4
)

// String returns the string representation of the compression type.
func (c CompressionType) String() string {
	switch c {
	case CompressionNone:
		return "none"
	case CompressionGzip:
		return "gzip"
	case CompressionSnappy:
		return "snappy"
	case CompressionLZ4:
		return "lz4"
	default:
		return "unknown"
	}
}

// Compressor handles message compression and decompression.
type Compressor struct {
	gzipWriterPool sync.Pool
	lz4WriterPool  sync.Pool
}

func NewCompressor() *Compressor {
	return &Compressor{
		gzipWriterPool: sync.Pool{
			New: func() interface{} {
				return gzip.NewWriter(nil)
			},
		},
		lz4WriterPool: sync.Pool{
			New: func() interface{} {
				return lz4.NewWriter(nil)
			},
		},
	}
}

// Compress compresses data using the specified algorithm.
func (c *Compressor) Compress(data []byte, compressionType CompressionType) ([]byte, error) {
	switch compressionType {
	case CompressionNone:
		return data, nil
	case CompressionGzip:
		return c.compressGzip(data)
	case CompressionSnappy:
		return c.compressSnappy(data)
	case CompressionLZ4:
		return c.compressLZ4(data)
	default:
		return nil, fmt.Errorf("unsupported compression type: %d", compressionType)
	}
}

// Decompress decompresses data using the specified algorithm.
func (c *Compressor) Decompress(data []byte, compressionType CompressionType) ([]byte, error) {
	switch compressionType {
	case CompressionNone:
		return data, nil
	case CompressionGzip:
		return c.decompressGzip(data)
	case CompressionSnappy:
		return c.decompressSnappy(data)
	case CompressionLZ4:
		return c.decompressLZ4(data)
	default:
		return nil, fmt.Errorf("unsupported compression type: %d", compressionType)
	}
}

func (c *Compressor) compressGzip(data []byte) ([]byte, error) {
	var buf bytes.Buffer

	w := c.gzipWriterPool.Get().(*gzip.Writer)
	w.Reset(&buf)

	if _, err := w.Write(data); err != nil {
		c.gzipWriterPool.Put(w)
		return nil, fmt.Errorf("gzip write: %w", err)
	}

	if err := w.Close(); err != nil {
		c.gzipWriterPool.Put(w)
		return nil, fmt.Errorf("gzip close: %w", err)
	}

	c.gzipWriterPool.Put(w)
	return buf.Bytes(), nil
}

func (c *Compressor) decompressGzip(data []byte) ([]byte, error) {
	r, err := gzip.NewReader(bytes.NewReader(data))
	if err != nil {
		return nil, fmt.Errorf("gzip reader: %w", err)
	}
	defer r.Close()

	result, err := io.ReadAll(r)
	if err != nil {
		return nil, fmt.Errorf("gzip read: %w", err)
	}

	return result, nil
}

// compressSnappy compresses data using Snappy algorithm.
// Snappy is optimized for speed over compression ratio.
func (c *Compressor) compressSnappy(data []byte) ([]byte, error) {
	// Snappy encode - no streaming needed for typical message sizes
	return snappy.Encode(nil, data), nil
}

// decompressSnappy decompresses Snappy-compressed data.
func (c *Compressor) decompressSnappy(data []byte) ([]byte, error) {
	decoded, err := snappy.Decode(nil, data)
	if err != nil {
		return nil, fmt.Errorf("snappy decode: %w", err)
	}
	return decoded, nil
}

// compressLZ4 compresses data using LZ4 algorithm.
// LZ4 provides a good balance of speed and compression ratio.
func (c *Compressor) compressLZ4(data []byte) ([]byte, error) {
	var buf bytes.Buffer

	sizeBuf := make([]byte, 4)
	binary.LittleEndian.PutUint32(sizeBuf, uint32(len(data)))
	buf.Write(sizeBuf)

	w := c.lz4WriterPool.Get().(*lz4.Writer)
	w.Reset(&buf)

	if _, err := w.Write(data); err != nil {
		c.lz4WriterPool.Put(w)
		return nil, fmt.Errorf("lz4 write: %w", err)
	}

	if err := w.Close(); err != nil {
		c.lz4WriterPool.Put(w)
		return nil, fmt.Errorf("lz4 close: %w", err)
	}

	c.lz4WriterPool.Put(w)
	return buf.Bytes(), nil
}

// decompressLZ4 decompresses LZ4-compressed data.
func (c *Compressor) decompressLZ4(data []byte) ([]byte, error) {
	if len(data) < 4 {
		return nil, fmt.Errorf("lz4 data too short")
	}

	uncompressedSize := binary.LittleEndian.Uint32(data[:4])

	r := lz4.NewReader(bytes.NewReader(data[4:]))

	result := make([]byte, uncompressedSize)
	n, err := io.ReadFull(r, result)
	if err != nil && err != io.EOF && err != io.ErrUnexpectedEOF {
		return nil, fmt.Errorf("lz4 read: %w", err)
	}

	return result[:n], nil
}

type CompressedRecord struct {
	Compression CompressionType
	Data        []byte
}

// WrapWithCompression compresses a record if needed.
func WrapWithCompression(data []byte, compression CompressionType) (*CompressedRecord, error) {
	compressor := NewCompressor()
	compressed, err := compressor.Compress(data, compression)
	if err != nil {
		return nil, err
	}

	return &CompressedRecord{
		Compression: compression,
		Data:        compressed,
	}, nil
}

// UnwrapCompression decompresses a record if needed.
func UnwrapCompression(record *CompressedRecord) ([]byte, error) {
	compressor := NewCompressor()
	return compressor.Decompress(record.Data, record.Compression)
}

// CompressionRatio calculates the compression ratio for the given data.
// Returns ratio as compressed_size/original_size (lower is better).
func CompressionRatio(original, compressed []byte) float64 {
	if len(original) == 0 {
		return 1.0
	}
	return float64(len(compressed)) / float64(len(original))
}