package ioz

import (
	"crypto/rand"
	"io"

	"google.golang.org/grpc/mem"
)

var (
	_ io.Reader      = (*stepReader)(nil)
	_ mem.BufferPool = (*testPool)(nil)
)

type readStep struct {
	n   int
	err error
}

type stepReader struct {
	reads []readStep
	read  []byte
}

func (s *stepReader) Read(buf []byte) (int, error) {
	if len(s.reads) == 0 {
		panic("unexpected Read() call")
	}
	read := s.reads[0]
	s.reads = s.reads[1:]
	_, err := rand.Read(buf[:read.n])
	if err != nil {
		panic(err)
	}
	s.read = append(s.read, buf[:read.n]...)
	return read.n, read.err
}

type testPool struct {
	allocated map[*[]byte]struct{}
}

func (t *testPool) Get(length int) *[]byte {
	buf := make([]byte, length)
	t.allocated[&buf] = struct{}{}
	return &buf
}

func (t *testPool) Put(buf *[]byte) {
	if _, ok := t.allocated[buf]; !ok {
		panic("unexpected put")
	}
	delete(t.allocated, buf)
}
