package zlog

import (
	"bytes"
	"log/slog"
	"sync"
)

type jsonHandlerPool struct {
	*sync.Pool
}

func newJsonHandlerPool(handlerOption *slog.HandlerOptions) *jsonHandlerPool {
	var handlerPool *jsonHandlerPool
	generator := &sync.Pool{
		New: func() any {
			buf := new(bytes.Buffer)
			return &jsonHandler{
				Buffer:      buf,
				pool:        handlerPool,
				JSONHandler: slog.NewJSONHandler(buf, handlerOption),
			}
		},
	}
	handlerPool = &jsonHandlerPool{
		generator,
	}
	return handlerPool
}

func (jhp *jsonHandlerPool) Get() *jsonHandler {
	handler := jhp.Pool.Get().(*jsonHandler)
	return handler
}

type jsonHandler struct {
	*bytes.Buffer
	pool *jsonHandlerPool
	*slog.JSONHandler
}

func (j *jsonHandler) Put() {
	// Only put back to pool if buffer is smaller than 4MB in RAM size.
	if j.Buffer.Cap() < 1024*1024*4 {
		j.Buffer.Reset()
		j.pool.Put(j)
	}
}

type bufferPool struct {
	*sync.Pool
}

func newBufferPool() *bufferPool {
	return &bufferPool{
		&sync.Pool{
			New: func() any {
				buf := new(bytes.Buffer)
				buf.Grow(1024)
				return buf
			},
		},
	}
}

func (bp *bufferPool) Get() *bytes.Buffer {
	return bp.Pool.Get().(*bytes.Buffer)
}

func (bp *bufferPool) Put(buf *bytes.Buffer) {
	if buf.Cap() < 1024*1024*4 {
		buf.Reset()
		bp.Pool.Put(buf)
	}
}