package udp

import (
	"bytes"
	"compress/zlib"
	"io"
)

func Compress(data []byte) ([]byte, error) {
	var buf bytes.Buffer

	// Write compression marker
	buf.WriteByte(0x5A)

	writer := zlib.NewWriter(&buf)
	_, err := writer.Write(data)
	if err != nil {
		return nil, err
	}

	err = writer.Close()
	if err != nil {
		return nil, err
	}

	return buf.Bytes(), nil
}

func Decompress(data []byte) ([]byte, error) {
	if len(data) == 0 {
		return data, nil
	}

	// Check compression marker
	if data[0] == 0xA5 {
		// Uncompressed data
		return data[1:], nil
	}

	if data[0] != 0x5A {
		// No compression marker, return as-is
		return data, nil
	}

	// Decompress zlib data
	reader := bytes.NewReader(data[1:])
	zlibReader, err := zlib.NewReader(reader)
	if err != nil {
		return nil, err
	}
	defer zlibReader.Close()

	var buf bytes.Buffer
	_, err = io.Copy(&buf, zlibReader)
	if err != nil {
		return nil, err
	}

	return buf.Bytes(), nil
}