package storage

import (
	"fmt"
	
	"github.com/klauspost/compress/zstd"
)

// CompressionService handles ZSTD compression and decompression
type CompressionService struct {
	compressor   *zstd.Encoder
	decompressor *zstd.Decoder
}

// NewCompressionService creates a new compression service
func NewCompressionService() (*CompressionService, error) {
	// Initialize ZSTD compressor
	compressor, err := zstd.NewWriter(nil)
	if err != nil {
		return nil, fmt.Errorf("failed to initialize ZSTD compressor: %v", err)
	}

	// Initialize ZSTD decompressor
	decompressor, err := zstd.NewReader(nil)
	if err != nil {
		compressor.Close()
		return nil, fmt.Errorf("failed to initialize ZSTD decompressor: %v", err)
	}

	return &CompressionService{
		compressor:   compressor,
		decompressor: decompressor,
	}, nil
}

// Close closes the compression and decompression resources
func (c *CompressionService) Close() {
	if c.compressor != nil {
		c.compressor.Close()
	}
	if c.decompressor != nil {
		c.decompressor.Close()
	}
}

// CompressData compresses data using ZSTD
func (c *CompressionService) CompressData(data []byte) ([]byte, error) {
	if c.compressor == nil {
		return nil, fmt.Errorf("compressor not initialized")
	}
	return c.compressor.EncodeAll(data, make([]byte, 0, len(data))), nil
}

// DecompressData decompresses ZSTD-compressed data
func (c *CompressionService) DecompressData(compressedData []byte) ([]byte, error) {
	if c.decompressor == nil {
		return nil, fmt.Errorf("decompressor not initialized")
	}
	return c.decompressor.DecodeAll(compressedData, nil)
}