package protocol

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

	"spark-server/internal/pkg/errs"
)

func Compress(body []byte) ([]byte, error) {
	var buf bytes.Buffer
	w := gzip.NewWriter(&buf)
	if _, err := w.Write(body); err != nil {
		_ = w.Close()
		return nil, err
	}
	if err := w.Close(); err != nil {
		return nil, err
	}
	return buf.Bytes(), nil
}

func Decompress(body []byte, limit int) ([]byte, error) {
	r, err := gzip.NewReader(bytes.NewReader(body))
	if err != nil {
		return nil, errs.ErrInvalidFrame
	}
	defer r.Close()
	out, err := io.ReadAll(io.LimitReader(r, int64(limit)+1))
	if err != nil {
		return nil, errs.ErrInvalidFrame
	}
	if len(out) > limit {
		return nil, errs.ErrFrameTooLarge
	}
	return out, nil
}

// PrepareBody gzip-compresses body when it is at least minCompress bytes and smaller after gzip.
func PrepareBody(body []byte, minCompress int) ([]byte, uint16, error) {
	if minCompress <= 0 || len(body) < minCompress {
		return body, 0, nil
	}
	out, err := Compress(body)
	if err != nil {
		return nil, 0, err
	}
	if len(out) >= len(body) {
		return body, 0, nil
	}
	return out, FlagCompress, nil
}

// DecodeBody inflates a gzip body when FlagCompress is set.
func DecodeBody(body []byte, flags uint16, limit int) ([]byte, error) {
	if flags&FlagCompress == 0 {
		return body, nil
	}
	return Decompress(body, limit)
}