package buffer_test

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

type pool struct {
	allocSize int
	alloc, free int
	errAfter []int
	zeroAfter []int
	shortAfter []int
	longAfter []int
	randomSize []int
}

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) 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 p.allocCondition(&p.shortAfter) {
		n /= 2
	}

	if p.allocCondition(&p.longAfter) {
		n *= 2
	}

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

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

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

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

func TestPoolUsage(t *testing.T) {
	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 := buffer.ReaderFrom(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 := buffer.ReaderFrom(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 := buffer.ReaderFrom(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 := buffer.ReaderFrom(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 := buffer.ReaderFrom(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 := buffer.ReaderFrom(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 := buffer.ReaderFrom(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 := buffer.ReaderFrom(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 := buffer.ReaderFrom(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 := buffer.ReaderFrom(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 := buffer.ReaderFrom(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 := buffer.ReaderFrom(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 := buffer.ReaderFrom(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 := buffer.ReaderFrom(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 := buffer.ReaderFrom(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 := buffer.ReaderFrom(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 := buffer.ReaderFrom(g, o)
			
			var b bytes.Buffer
			_, err := r.WriteTo(&b)
			if !errors.Is(err, buffer.ErrZeroAllocation) {
				t.Fatal("failed to fail", err)
			}
		})
	})

	// varying segment sizes: read bytes, peek
	// one byte segments: read, read bytes, read utf8, peek, write to
	// pool error on allocate: read, read bytes, read utf8, peek, write to
}