package pool

import (
	"sync"
)

// Buffer contains a byte slice.
type Buffer struct {
	Data []byte
}

// BufferPool manages a pool of Buffers.
type BufferPool struct {
	poolSliceSize uint32 // Size for the buffer slices in the pool.
	buffersPool   *sync.Pool
}

// NewBufferPool creates a BufferPool with a specified size.
func NewBufferPool(poolSliceSize uint32) BufferPool {
	pool := sync.Pool{
		New: func() any {
			return new(Buffer)
		},
	}
	return BufferPool{poolSliceSize: poolSliceSize, buffersPool: &pool}
}

// Get retrieves a Buffer from the pool or creates a new one if necessary.
// It ensures the buffer's capacity is at least the specified size.
func (pool BufferPool) Get(size uint32) *Buffer {
	result := pool.buffersPool.Get().(*Buffer)

	if cap(result.Data) < int(size) {
		result.Data = make([]byte, size)
	} else {
		result.Data = result.Data[:size]
	}
	return result
}

// Put returns a Buffer to the pool if its capacity does not exceed poolSliceSize.
func (pool BufferPool) Put(buf *Buffer) {
	if cap(buf.Data) > int(pool.poolSliceSize) {
		return
	}

	buf.Data = buf.Data[:0]
	pool.buffersPool.Put(buf)
}

// PoolSliceSize returns the size for buffer slices in the pool.
func (pool BufferPool) PoolSliceSize() uint32 {
	return uint32(pool.poolSliceSize)
}