package buffer

import (
	"encoding/binary"
	"errors"
	"io"
)

// MaxUint16 is the maximum value of a uint16. It is used to check if the
// length of the data is too large to be encoded in a uint16.
const MaxUint16 = 0xFFFF

// ErrInvalidLength is the error message returned when the length of the data
// is too large to be encoded in a uint16.
var ErrInvalidLength = errors.New("data length is too large to be encoded in a uint16")

// NewLenPrefixed returns a new buffer with a length prefix. The length prefix is
// 2 bytes long and is encoded in big-endian order. The length prefix is
// followed by the data.
func NewLenPrefixed(data []byte) ([]byte, error) {
	length := len(data)
	// Overflow Guard: If the length of the data is greater than the maximum
	// value of a uint16, return an error.
	if length > MaxUint16 {
		return nil, ErrInvalidLength
	}
	buf := make([]byte, 2+len(data))
	binary.BigEndian.PutUint16(buf, uint16(len(data)))
	copy(buf[2:], data)
	return buf, nil
}

func ReadLenPrefixed(maxSize uint16, r io.Reader) ([]byte, error) {
	lenBuf := make([]byte, 2)
	if _, err := io.ReadFull(r, lenBuf); err != nil {
		return nil, err
	}
	if len(lenBuf) < 2 {
		return nil, errors.New("buffer is too small to contain length prefix")
	}
	length := binary.BigEndian.Uint16(lenBuf)

	if length > maxSize {
		return nil, errors.New("data length is too large to be encoded in a uint16")
	}

	buf := make([]byte, length)
	if _, err := io.ReadFull(r, buf); err != nil {
		return nil, err
	}

	if len(buf) < int(length) {
		return nil, errors.New("buffer is too small to contain data")
	}

	return buf, nil
}