package pdu

import (
	"bytes"
	"log"
	"sync"
)

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

func (bpm *BufferPoolManager) logf(format string, args ...interface{}) {
	if bpm.debug {
		log.Printf(format, args...)
	}
}

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 {
			bpm.logf("Creating new pool for size %d\n", size)
			pool = &sync.Pool{
				New: func() interface{} {
					bpm.logf("Creating new buffer of size %d\n", size)
					buf := bytes.NewBuffer(make([]byte, size))
					buf.Reset()
					return buf
				},
			}
			bpm.pools[size] = pool
		}
		bpm.mu.Unlock()
	}

	buf := pool.Get().(*bytes.Buffer)
	bpm.logf("Returning buffer of size %d: %p\n", buf.Cap(), buf)
	return buf
}

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

	if !exists {
		bpm.logf("Cannot return %p, No pool for size %d\n", buf, size)
		return
	}

	buf.Reset()
	pool.Put(buf)
	bpm.logf("Returned buffer of size %d: %p\n", size, buf)
}