package buffer_test

import (
	"bytes"
	"code.squareroundforest.org/arpio/buffer"
	"errors"
	"io"
	"math/rand"
	"testing"
)

type pool struct {
	allocSize   int
	alloc, free int
	errAfter    []int
	zeroAfter   []int
	varyingSize []int
	rand        *rand.Rand
}

func (p pool) allocCondition(c *[]int) bool {
	if len(*c) == 0 {
		return false
	}

	if p.alloc < (*c)[0] {
		return false
	}

	*c = (*c)[1:]
	return true
}

func (p *pool) ensureRand() {
	if p.rand != nil {
		return
	}

	p.rand = rand.New(rand.NewSource(9))
}

func (p *pool) Get() ([]byte, error) {
	defer func() {
		p.alloc++
	}()

	if p.allocCondition(&p.errAfter) {
		return nil, errTest
	}

	if p.allocCondition(&p.zeroAfter) {
		return nil, nil
	}

	n := p.allocSize
	if len(p.varyingSize) > 1 {
		p.ensureRand()
		n = p.varyingSize[0] + p.rand.Intn(p.varyingSize[1])
	}

	if len(p.varyingSize) == 1 {
		p.ensureRand()
		n = 1 + p.rand.Intn(p.varyingSize[0])
	}

	return make([]byte, n), nil
}

func (p *pool) Put([]byte) {
	p.free++
}

func TestPoolUsage(t *testing.T) {
	for _, cr := range []createReader{buffer.BufferedReader, testContent} {
		t.Run("allocate", func(t *testing.T) {
			t.Run("read", func(t *testing.T) {
				g := &gen{max: 1 << 15}
				p := &pool{allocSize: 1 << 12}
				o := buffer.Options{Pool: p}
				r := cr(g, o)
				b := make([]byte, 256)
				n, err := r.Read(b)
				if err != nil {
					t.Fatal(err)
				}

				if n != len(b) {
					t.Fatal("invalid read length")
				}

				if !bytes.Equal(b, generate(len(b))) {
					t.Fatal("invalid content")
				}

				if p.alloc != 1 {
					t.Fatal("invalid allocation count")
				}
			})

			t.Run("read bytes", func(t *testing.T) {
				g := &gen{max: 1 << 15}
				p := &pool{allocSize: 1 << 9}
				o := buffer.Options{Pool: p}
				r := cr(g, o)
				b, ok, err := r.ReadBytes([]byte("123"), 1<<12)
				if err != nil {
					t.Fatal(err)
				}

				if ok {
					t.Fatal("unexpected delimiter")
				}

				if len(b) != 0 {
					t.Fatal("invalid content")
				}

				if p.alloc != 8 {
					t.Fatal("invalid allocation count")
				}
			})

			t.Run("read utf8", func(t *testing.T) {
				g := &gen{max: 1 << 15}
				p := &pool{allocSize: 1 << 12}
				o := buffer.Options{Pool: p}
				r := cr(g, o)
				runes, n, err := r.ReadUTF8(1 << 12)
				if err != nil {
					t.Fatal(err)
				}

				if n != 1<<12 {
					t.Fatal("unexpected delimiter")
				}

				if len(runes) != 1<<12 {
					t.Fatal("invalid content")
				}

				if p.alloc != 3 {
					t.Fatal("invalid allocation count")
				}
			})

			t.Run("peek", func(t *testing.T) {
				g := &gen{max: 1 << 15}
				p := &pool{allocSize: 1 << 12}
				o := buffer.Options{Pool: p}
				r := cr(g, o)
				b, err := r.Peek(3 * 1 << 12)
				if err != nil {
					t.Fatal(err)
				}

				if len(b) != 3*1<<12 {
					t.Fatal("invalid content")
				}

				if p.alloc != 3 {
					t.Fatal("invalid allocation count")
				}
			})

			t.Run("buffered", func(t *testing.T) {
				g := &gen{max: 1 << 15}
				p := &pool{allocSize: 1 << 12}
				o := buffer.Options{Pool: p}
				r := cr(g, o)
				b := r.Buffered()
				if len(b) != 0 {
					t.Fatal("invalid content")
				}

				if p.alloc != 0 {
					t.Fatal("invalid allocation count")
				}
			})

			t.Run("write to", func(t *testing.T) {
				g := &gen{max: 1 << 15}
				p := &pool{allocSize: 1 << 12}
				o := buffer.Options{Pool: p}
				r := cr(g, o)

				var b bytes.Buffer
				n, err := r.WriteTo(&b)
				if err != nil {
					t.Fatal(err)
				}

				if n != 1<<15 {
					t.Fatal("invalid write length")
				}

				if !bytes.Equal(b.Bytes(), generate(1<<15)) {
					t.Fatal("invalid content")
				}

				if p.alloc != 1 {
					t.Fatal("invalid allocation count")
				}
			})
		})

		t.Run("free", func(t *testing.T) {
			t.Run("read", func(t *testing.T) {
				g := &gen{max: 1 << 15}
				p := &pool{allocSize: 1 << 12}
				o := buffer.Options{Pool: p}
				r := cr(g, o)
				b := make([]byte, 1<<9)
				for {
					_, err := r.Read(b)
					if errors.Is(err, io.EOF) {
						break
					}

					if err != nil {
						t.Fatal(err)
					}
				}

				if p.alloc != 1 {
					t.Fatal("invalid allocation count")
				}

				if p.free != 1 {
					t.Fatal("invalid free count")
				}
			})

			t.Run("read bytes", func(t *testing.T) {
				g := &gen{max: 1 << 15}
				p := &pool{allocSize: 1 << 12}
				o := buffer.Options{Pool: p}
				r := cr(g, o)
				_, _, err := r.ReadBytes([]byte("123"), 1<<15+3)
				if err != nil {
					t.Fatal(err)
				}

				_, _, err = r.ReadBytes([]byte("123"), 1<<15+3)
				if !errors.Is(err, io.EOF) {
					t.Fatal(err)
				}

				if p.alloc != 9 {
					t.Fatal("invalid allocation count", p.alloc)
				}

				if p.free != 9 {
					t.Fatal("invalid free count", p.free)
				}
			})

			t.Run("read utf8", func(t *testing.T) {
				g := &gen{max: 1 << 15}
				p := &pool{allocSize: 1 << 12}
				o := buffer.Options{Pool: p}
				r := cr(g, o)
				for {
					runes, n, err := r.ReadUTF8(1 << 12)
					if errors.Is(err, io.EOF) {
						break
					}

					if err != nil {
						t.Fatal(err)
					}

					if n != 1<<12 {
						t.Fatal("unexpected delimiter")
					}

					if len(runes) != 1<<12 {
						t.Fatal("invalid content")
					}
				}

				if p.alloc != p.free {
					t.Fatal("invalid allocation count", p.alloc, p.free)
				}
			})

			t.Run("peek", func(t *testing.T) {
				g := &gen{max: 1 << 15}
				p := &pool{allocSize: 1 << 12}
				o := buffer.Options{Pool: p}
				r := cr(g, o)
				b, err := r.Peek(3 * 1 << 12)
				if err != nil {
					t.Fatal(err)
				}

				if len(b) != 3*1<<12 {
					t.Fatal("invalid content")
				}

				if p.alloc != 3 {
					t.Fatal("invalid allocation count")
				}

				if p.free != 0 {
					t.Fatal("invalid allocation count")
				}

				b = make([]byte, 1<<9)
				for {
					_, err := r.Read(b)
					if errors.Is(err, io.EOF) {
						break
					}

					if err != nil {
						t.Fatal(err)
					}
				}

				if p.alloc != 3 {
					t.Fatal("invalid allocation count")
				}

				if p.free != 3 {
					t.Fatal("invalid allocation count")
				}
			})

			t.Run("buffered", func(t *testing.T) {
				g := &gen{max: 1 << 15}
				p := &pool{allocSize: 1 << 12}
				o := buffer.Options{Pool: p}
				r := cr(g, o)
				b := r.Buffered()
				if len(b) != 0 {
					t.Fatal("invalid content")
				}

				if p.alloc != 0 || p.free != 0 {
					t.Fatal("invalid allocation count")
				}
			})

			t.Run("write to", func(t *testing.T) {
				g := &gen{max: 1 << 15}
				p := &pool{allocSize: 1 << 12}
				o := buffer.Options{Pool: p}
				r := cr(g, o)

				var b bytes.Buffer
				n, err := r.WriteTo(&b)
				if err != nil {
					t.Fatal(err)
				}

				if n != 1<<15 {
					t.Fatal("invalid write length")
				}

				if !bytes.Equal(b.Bytes(), generate(1<<15)) {
					t.Fatal("invalid content")
				}

				if p.alloc != 1 {
					t.Fatal("invalid allocation count")
				}

				if p.free != 1 {
					t.Fatal("invalid free count")
				}
			})
		})

		t.Run("null segment", func(t *testing.T) {
			t.Run("read", func(t *testing.T) {
				g := &gen{max: 1 << 15}
				p := &pool{
					allocSize: 1 << 12,
					zeroAfter: []int{0},
				}

				o := buffer.Options{Pool: p}
				r := cr(g, o)
				b := make([]byte, 256)
				_, err := r.Read(b)
				if !errors.Is(err, buffer.ErrZeroAllocation) {
					t.Fatal("failed to fail", err)
				}
			})

			t.Run("read bytes", func(t *testing.T) {
				g := &gen{max: 1 << 15}
				p := &pool{
					allocSize: 1 << 9,
					zeroAfter: []int{1},
				}

				o := buffer.Options{Pool: p}
				r := cr(g, o)
				_, _, err := r.ReadBytes([]byte("123"), 1<<12)
				if err != nil {
					t.Fatal(err)
				}

				_, _, err = r.ReadBytes([]byte("123"), 1<<12)
				if !errors.Is(err, buffer.ErrZeroAllocation) {
					t.Fatal("failed to fail", err)
				}
			})

			t.Run("read utf8", func(t *testing.T) {
				g := &gen{max: 1 << 15}
				p := &pool{
					allocSize: 1 << 12,
					zeroAfter: []int{0},
				}

				o := buffer.Options{Pool: p}
				r := cr(g, o)
				_, _, err := r.ReadUTF8(1 << 12)
				if !errors.Is(err, buffer.ErrZeroAllocation) {
					t.Fatal("failed to fail", err)
				}
			})

			t.Run("peek", func(t *testing.T) {
				g := &gen{max: 1 << 15}
				p := &pool{
					allocSize: 1 << 12,
					zeroAfter: []int{0},
				}

				o := buffer.Options{Pool: p}
				r := cr(g, o)
				_, err := r.Peek(3 * 1 << 12)
				if !errors.Is(err, buffer.ErrZeroAllocation) {
					t.Fatal("failed to fail", err)
				}
			})

			t.Run("write to", func(t *testing.T) {
				g := &gen{max: 1 << 15}
				p := &pool{
					allocSize: 1 << 12,
					zeroAfter: []int{0},
				}

				o := buffer.Options{Pool: p}
				r := cr(g, o)

				var b bytes.Buffer
				_, err := r.WriteTo(&b)
				if !errors.Is(err, buffer.ErrZeroAllocation) {
					t.Fatal("failed to fail", err)
				}
			})
		})

		t.Run("varying segment sizes", func(t *testing.T) {
			t.Run("read bytes", func(t *testing.T) {
				t.Run("find", func(t *testing.T) {
					g := &gen{
						max:                1 << 15,
						customContentAfter: []int{1 << 11},
						customContent:      map[int][]byte{1 << 11: []byte("123")},
					}

					p := &pool{varyingSize: []int{8, 256}}
					o := buffer.Options{Pool: p}
					r := cr(g, o)
					b, ok, err := r.ReadBytes([]byte("123"), 1<<15)
					if err != nil {
						t.Fatal(err)
					}

					if !ok {
						t.Fatal("failed to find delimiter")
					}

					if !bytes.Equal(b, append(generate(1<<11), []byte("123")...)) {
						t.Fatal("invalid content")
					}
				})

				t.Run("find not", func(t *testing.T) {
					g := &gen{max: 1 << 15}
					p := &pool{varyingSize: []int{8, 256}}
					o := buffer.Options{Pool: p}
					r := cr(g, o)
					b, ok, err := r.ReadBytes([]byte("123"), 1<<15)
					if err != nil {
						t.Fatal(err)
					}

					if ok {
						t.Fatal("invalid delimiter")
					}

					if !bytes.Equal(b, generate(1<<15)) {
						t.Fatal("invalid content")
					}
				})
			})

			t.Run("peek", func(t *testing.T) {
				g := &gen{max: 1 << 15}
				p := &pool{varyingSize: []int{8, 256}}
				o := buffer.Options{Pool: p}
				r := cr(g, o)
				b, err := r.Peek(1 << 11)
				if err != nil {
					t.Fatal(err)
				}

				if !bytes.Equal(b, generate(1<<11)) {
					t.Fatal("invalid content")
				}
			})
		})

		t.Run("one byte segments", func(t *testing.T) {
			t.Run("read", func(t *testing.T) {
				g := &gen{max: 1 << 15}
				p := &pool{allocSize: 1}
				o := buffer.Options{Pool: p}
				r := cr(g, o)
				b, err := io.ReadAll(r)
				if err != nil {
					t.Fatal(err)
				}

				if !bytes.Equal(b, generate(1<<15)) {
					t.Fatal("invalid content")
				}
			})

			t.Run("read bytes", func(t *testing.T) {
				t.Run("find", func(t *testing.T) {
					g := &gen{
						max:                1 << 15,
						customContentAfter: []int{1 << 11},
						customContent:      map[int][]byte{1 << 11: []byte("123")},
					}

					p := &pool{allocSize: 1}
					o := buffer.Options{Pool: p}
					r := cr(g, o)
					b, ok, err := r.ReadBytes([]byte("123"), 1<<15)
					if err != nil {
						t.Fatal(err)
					}

					if !ok {
						t.Fatal("failed to find delimiter")
					}

					if !bytes.Equal(b, append(generate(1<<11), []byte("123")...)) {
						t.Fatal("invalid content")
					}
				})

				t.Run("find not", func(t *testing.T) {
					g := &gen{max: 1 << 15}
					p := &pool{allocSize: 1}
					o := buffer.Options{Pool: p}
					r := cr(g, o)
					b, ok, err := r.ReadBytes([]byte("123"), 1<<15)
					if err != nil {
						t.Fatal(err)
					}

					if ok {
						t.Fatal("invalid delimiter found")
					}

					if len(b) != 0 {
						t.Fatal("invalid content", len(b))
					}
				})
			})

			t.Run("read utf8", func(t *testing.T) {
				g := &gen{
					max: 1 << 15,
					rng: utf8W2Range,
				}

				p := &pool{allocSize: 1}
				o := buffer.Options{Pool: p}
				r := cr(g, o)
				runes, n, err := r.ReadUTF8(1 << 14)
				if err != nil {
					t.Fatal(err)
				}

				if n != 1<<15 {
					t.Fatal("invalid read length")
				}

				if string(runes) != string(generateFrom(utf8W2Range, 1<<15)) {
					t.Fatal("invalid content")
				}
			})

			t.Run("peek", func(t *testing.T) {
				g := &gen{max: 1 << 15}
				p := &pool{allocSize: 1}
				o := buffer.Options{Pool: p}
				r := cr(g, o)
				b, err := r.Peek(1 << 14)
				if err != nil {
					t.Fatal(err)
				}

				if !bytes.Equal(b, generate(1<<14)) {
					t.Fatal("invalid content")
				}
			})

			t.Run("write to", func(t *testing.T) {
				g := &gen{max: 1 << 15}
				p := &pool{allocSize: 1}
				o := buffer.Options{Pool: p}
				r := cr(g, o)

				var b bytes.Buffer
				n, err := r.WriteTo(&b)
				if err != nil {
					t.Fatal(err)
				}

				if n != 1<<15 {
					t.Fatal("invalid write length")
				}

				if !bytes.Equal(b.Bytes(), generate(1<<15)) {
					t.Fatal("invalid content")
				}
			})
		})

		t.Run("pool error on allocate", func(t *testing.T) {
			t.Run("read", func(t *testing.T) {
				g := &gen{max: 1 << 15}
				p := &pool{
					allocSize: 1 << 11,
					errAfter:  []int{0},
				}

				o := buffer.Options{Pool: p}
				r := cr(g, o)
				b, err := io.ReadAll(r)
				if !errors.Is(err, errTest) {
					t.Fatal("failed to fail")
				}

				if len(b) != 0 {
					t.Fatal("invalid content")
				}
			})

			t.Run("read bytes", func(t *testing.T) {
				t.Run("immediate", func(t *testing.T) {
					g := &gen{max: 1 << 15}
					p := &pool{
						allocSize: 1 << 11,
						errAfter:  []int{0},
					}

					o := buffer.Options{Pool: p}
					r := cr(g, o)
					b, ok, err := r.ReadBytes([]byte("123"), 1<<12)
					if !errors.Is(err, errTest) {
						t.Fatal("failed to fail with the right error", err)
					}

					if ok {
						t.Fatal("invalid delimiter found")
					}

					if len(b) != 0 {
						t.Fatal("invalid content")
					}
				})

				t.Run("on grow", func(t *testing.T) {
					g := &gen{max: 1 << 15}
					p := &pool{
						allocSize: 1 << 11,
						errAfter:  []int{1},
					}

					o := buffer.Options{Pool: p}
					r := cr(g, o)
					b, ok, err := r.ReadBytes([]byte("123"), 1<<15)
					if err != nil {
						t.Fatal(err)
					}

					if ok {
						t.Fatal("invalid delimiter found")
					}

					if !bytes.Equal(b, generate(1<<11)) {
						t.Fatal("invalid content")
					}

					b, ok, err = r.ReadBytes([]byte("123"), 1<<15)
					if !errors.Is(err, errTest) {
						t.Fatal("failed to fail with the right error", err)
					}

					if ok {
						t.Fatal("invalid delimiter found")
					}

					if len(b) != 0 {
						t.Fatal("invalid content")
					}
				})
			})

			t.Run("read utf8", func(t *testing.T) {
				g := &gen{
					max: 1 << 15,
					rng: utf8W2Range,
				}

				p := &pool{
					allocSize: 1 << 11,
					errAfter:  []int{0},
				}

				o := buffer.Options{Pool: p}
				r := cr(g, o)
				runes, n, err := r.ReadUTF8(1 << 14)
				if !errors.Is(err, errTest) {
					t.Fatal(err)
				}

				if n != 0 {
					t.Fatal("invalid read length")
				}

				if len(runes) != 0 {
					t.Fatal("invalid content")
				}
			})

			t.Run("peek", func(t *testing.T) {
				t.Run("immediate", func(t *testing.T) {
					g := &gen{max: 1 << 15}
					p := &pool{
						allocSize: 1 << 11,
						errAfter:  []int{0},
					}

					o := buffer.Options{Pool: p}
					r := cr(g, o)
					b, err := r.Peek(1 << 12)
					if !errors.Is(err, errTest) {
						t.Fatal("failed to fail with the right error", err)
					}

					if len(b) != 0 {
						t.Fatal("invalid content")
					}
				})

				t.Run("on grow", func(t *testing.T) {
					g := &gen{max: 1 << 15}
					p := &pool{
						allocSize: 1 << 11,
						errAfter:  []int{1},
					}

					o := buffer.Options{Pool: p}
					r := cr(g, o)
					b, err := r.Peek(1 << 15)
					if err != nil {
						t.Fatal(err)
					}

					if !bytes.Equal(b, generate(1<<11)) {
						t.Fatal("invalid content")
					}

					b, ok, err := r.ReadBytes([]byte("123"), 1<<11)
					if err != nil {
						t.Fatal(err)
					}

					if ok {
						t.Fatal("invalid delimiter found")
					}

					if !bytes.Equal(b, generate(1<<11)) {
						t.Fatal("invalid content")
					}

					b, err = r.Peek(1 << 15)
					if !errors.Is(err, errTest) {
						t.Fatal("failed to fail with the right error")
					}

					if len(b) != 0 {
						t.Fatal("invalid content")
					}
				})
			})

			t.Run("write to", func(t *testing.T) {
				g := &gen{max: 1 << 15}
				p := &pool{
					allocSize: 1 << 11,
					errAfter:  []int{0},
				}

				o := buffer.Options{Pool: p}
				r := cr(g, o)

				var b bytes.Buffer
				n, err := r.WriteTo(&b)
				if !errors.Is(err, errTest) {
					t.Fatal("failed to fail with the right error")
				}

				if n != 0 {
					t.Fatal("invalid write length")
				}

				if b.Len() != 0 {
					t.Fatal("invalid content")
				}
			})
		})
	}
}

func TestDefaultPool(t *testing.T) {
	g := &gen{max: 1 << 18}
	r := buffer.BufferedReader(g, buffer.Options{})
	b, err := io.ReadAll(r)
	if err != nil {
		t.Fatal(err)
	}

	if !bytes.Equal(b, generate(1<<18)) {
		t.Fatal("output does not match", len(b))
	}
}