package pdu

import (
	"sync"
)

type BufferPoolManager struct {
	pools map[uint]*sync.Pool
	mu    sync.RWMutex
}

func NewBufferPoolManager() *BufferPoolManager {
	return &BufferPoolManager{
		pools: make(map[uint]*sync.Pool),
	}
}

func (bpm *BufferPoolManager) Get(size uint) *[]uint8 {
	bpm.mu.RLock()
	pool, exists := bpm.pools[size]
	bpm.mu.RUnlock()

	if !exists {
		bpm.mu.Lock()
		// Double-check if another goroutine added the pool while we were waiting
		pool, exists = bpm.pools[size]
		if !exists {
			pool = &sync.Pool{
				New: func() interface{} {
					buf := make([]uint8, size)
					return &buf
				},
			}
			bpm.pools[size] = pool
		}
		bpm.mu.Unlock()
	}

	return pool.Get().(*[]uint8)
}

func (bpm *BufferPoolManager) Put(buf *[]uint8) {
	size := uint(len(*buf))
	bpm.mu.RLock()
	pool, exists := bpm.pools[size]
	bpm.mu.RUnlock()

	if !exists {
		return
	}

	// Clear buffer
	for i := range *buf {
		(*buf)[i] = 0
	}

	pool.Put(buf)
}