package pdu

import (
	"bytes"
	"sync"
)

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

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

func (bpm *BufferPoolManager) Get(size int) *bytes.Buffer {
	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{} {
					return bytes.NewBuffer(make([]byte, size))
				},
			}
			bpm.pools[size] = pool
		}
		bpm.mu.Unlock()
	}

	return pool.Get().(*bytes.Buffer);
}

func (bpm *BufferPoolManager) Put(buf *bytes.Buffer) {
	size := buf.Len()
	bpm.mu.RLock()
	pool, exists := bpm.pools[size]
	bpm.mu.RUnlock()

	if !exists {
		return
	}

	buf.Reset()
	pool.Put(buf)
}