package utils

import (
	"io"
	"sync"

	"braces.dev/errtrace"
	"github.com/tidbcloud/tiinf/monitor/tilog"
	"go.uber.org/zap"
)

// ByteBuffer implements a simple byte buffer.
type ByteBuffer struct {
	// B is the underlying byte slice.
	B []byte
}

// Reset resets bb.
func (bb *ByteBuffer) Reset() {
	bb.B = bb.B[:0]
}

// Write appends p to bb.
func (bb *ByteBuffer) Write(p []byte) (int, error) {
	bb.B = append(bb.B, p...)
	return len(p), nil
}

// NewReader returns new reader for the given bb.
func (bb *ByteBuffer) NewReader() io.Reader {
	return &reader{
		bb: bb,
	}
}

type reader struct {
	bb *ByteBuffer

	// readOffset is the offset in bb.B for read.
	readOffset int
}

// Read reads up to len(p) bytes from bb.
func (r *reader) Read(p []byte) (int, error) {
	var err error
	n := copy(p, r.bb.B[r.readOffset:])
	if n < len(p) {
		err = io.EOF
	}
	r.readOffset += n
	return n, errtrace.Wrap(err)
}

// ByteBufferPool is a pool of ByteBuffers.
type ByteBufferPool struct {
	p sync.Pool
}

// Get obtains a ByteBuffer from bbp.
func (bbp *ByteBufferPool) Get() *ByteBuffer {
	bbv := bbp.p.Get()
	if bbv == nil {
		return &ByteBuffer{}
	}
	bf, ok := bbv.(*ByteBuffer)
	if !ok {
		tilog.PlainError("ByteBufferPool.Get: unexpected type", zap.Any("bbv", bbv))
	}
	return bf
}

// Put puts bb into bbp.
func (bbp *ByteBufferPool) Put(bb *ByteBuffer) {
	bb.Reset()
	bbp.p.Put(bb)
}