package ws

import (
	"bytes"
	"compress/gzip"
	"io"
	"sync"
)

const gzipMagicByte = 0x1F

// Compressor 管理 gzip Writer/Reader 的 sync.Pool 复用
type Compressor struct {
	writerPool sync.Pool
	readerPool sync.Pool
	level      int
}

// NewCompressor 创建 Compressor（level: gzip.DefaultCompression/BestSpeed/BestCompression 等）
func NewCompressor(level int) *Compressor {
	return &Compressor{
		level: level,
		writerPool: sync.Pool{
			New: func() any {
				w, _ := gzip.NewWriterLevel(io.Discard, level)
				return w
			},
		},
		readerPool: sync.Pool{
			New: func() any { return new(gzip.Reader) },
		},
	}
}

// Compress 压缩数据，返回 gzip 字节
func (c *Compressor) Compress(data []byte) ([]byte, error) {
	buf := &bytes.Buffer{}
	w := c.writerPool.Get().(*gzip.Writer)
	w.Reset(buf)
	if _, err := w.Write(data); err != nil {
		c.writerPool.Put(w)
		return nil, err
	}
	if err := w.Close(); err != nil {
		c.writerPool.Put(w)
		return nil, err
	}
	c.writerPool.Put(w)
	return buf.Bytes(), nil
}

// Decompress 解压 gzip 数据
func (c *Compressor) Decompress(data []byte) ([]byte, error) {
	r := c.readerPool.Get().(*gzip.Reader)
	if err := r.Reset(bytes.NewReader(data)); err != nil {
		c.readerPool.Put(r)
		return nil, err
	}
	out, err := io.ReadAll(r)
	r.Close()
	c.readerPool.Put(r)
	return out, err
}

// IsCompressed 检查数据是否以 gzip magic byte 开头
func IsCompressed(data []byte) bool {
	return len(data) > 0 && data[0] == gzipMagicByte
}

// DefaultCompressor 全局默认压缩器（gzip.DefaultCompression = level 6）
var DefaultCompressor = NewCompressor(gzip.DefaultCompression)