package xturbo

import (
	"bytes"
	"encoding/binary"
	"errors"
	"syscall"
	"testing"

	"golang.org/x/sys/unix"
)

func TestBufRequest(t *testing.T) {
	b := NewBuf(4)
	p, seq := b.Request(8, 2, 4)
	if seq != 1 || len(p) != 4 || !bytes.Equal(b.Bytes(), []byte{8, 2, 2, 0, 0, 0, 0, 0}) {
		t.Fatalf("%x %d", b.Bytes(), seq)
	}
	p[0] = 255
	b.Reset()
	p, seq = b.Request(14, 0, 4)
	if seq != 2 || p[0] != 0 || b.Next() != 3 || b.Len() != 8 {
		t.Fatal("reset sequence/zeroing")
	}
	first := errors.New("first")
	b.Fail(first)
	b.Fail(ErrBadLength)
	p, seq = b.Request(1, 0, 0)
	if p != nil || seq != 0 || b.Err() != first {
		t.Fatal("sticky error")
	}
}

func TestNewBufDefault(t *testing.T) {
	b := NewBuf(0)
	if cap(b.b) != 64<<10 || b.maxReq != 262140 || b.Next() != 1 || b.Len() != 0 {
		t.Fatalf("cap %d maxReq %d next %d", cap(b.b), b.maxReq, b.Next())
	}
}

func TestBufLimits(t *testing.T) {
	for _, body := range []int{-1, 1, 262140, 1 << 20} {
		b := NewBuf(0)
		p, s := b.Request(1, 0, body)
		if p != nil || s != 0 || b.Err() == nil || b.Len() != 0 || b.Next() != 1 {
			t.Fatal(body)
		}
	}
	// Largest core request: 0xFFFF units, 4-byte header.
	b := NewBuf(0)
	if p, s := b.Request(1, 0, 262136); s != 1 || len(p) != 262136 || binary.LittleEndian.Uint16(b.Bytes()[2:]) != 0xffff {
		t.Fatal("largest core request")
	}
}

func TestBufMaxReq(t *testing.T) {
	b := NewBuf(0)
	b.setMaxReqWords(70000) // 280000 bytes
	if b.maxReq != 280000 {
		t.Fatal(b.maxReq)
	}
	// body+8 == 280000 fits exactly.
	if p, s := b.Request(1, 0, 280000-8); s != 1 || len(p) != 280000-8 {
		t.Fatal("max-size big request")
	}
	if got := binary.LittleEndian.Uint32(b.Bytes()[4:]); got != 70000 {
		t.Fatal(got)
	}
	if p, s := b.Request(1, 0, 280000-4); p != nil || s != 0 || b.Err() != ErrRequestTooLarge {
		t.Fatal("over maxReq accepted")
	}
	b = NewBuf(0)
	b.setMaxReqWords(0xffffffff)
	if b.maxReq != 0xfffffffc {
		t.Fatal("maxReq clamp", b.maxReq)
	}
}

func TestBufSeqOverflow(t *testing.T) {
	b := NewBuf(0)
	b.next = 0xffff
	if p, s := b.Request(1, 0, 0); p != nil || s != 0 || b.Err() != ErrSeqOverflow {
		t.Fatal("overflow")
	}
	b = NewBuf(0)
	b.next = 0xfffe
	if _, s := b.Request(1, 0, 0); s != 0xfffe {
		t.Fatal("last allowed seq", s)
	}
	if _, s := b.Request(1, 0, 0); s != 0 || b.Err() != ErrSeqOverflow {
		t.Fatal("overflow at next-acked == 0xffff")
	}
	b = NewBuf(0)
	b.next = 0xffff
	b.ack(5)
	b.ack(3) // stale acks never move acked backwards
	b.ack(0xffff)
	if b.acked != 5 {
		t.Fatal("ack", b.acked)
	}
	if _, s := b.Request(1, 0, 0); s != 0xffff || b.Err() != nil {
		t.Fatal("ack must free sequence space")
	}
}

func TestBufGrowth(t *testing.T) {
	b := NewBuf(16)
	for i := range 15 {
		b.Request(1, 0, 4)
		want := 16
		for want < (i+1)*8 {
			want *= 2
		}
		if c := cap(b.b); c != want {
			t.Fatal("cap", i, c, want)
		}
	}
	if cap(b.b) != 128 || b.Len() != 120 {
		t.Fatal(cap(b.b), b.Len())
	}
	// One request bigger than double the capacity grows straight to fit.
	b.Request(1, 0, 1000)
	if cap(b.b) != 2048 || b.Len() != 1124 {
		t.Fatal(cap(b.b), b.Len())
	}
	for i, v := range b.Bytes()[:120] {
		if i%8 == 0 && v != 1 || i%8 == 2 && v != 2 {
			t.Fatal("growth lost bytes", i)
		}
	}
	b.Reset()
	if cap(b.b) != 2048 {
		t.Fatal("shrunk")
	}
}

func TestAllocBufRequest(t *testing.T) {
	for _, body := range []int{4, 65536} {
		b := NewBuf(128 << 10)
		if n := testing.AllocsPerRun(1000, func() { b.Reset(); b.Request(1, 0, body) }); n != 0 {
			t.Fatal(n)
		}
		if b.Err() != nil {
			t.Fatal(b.Err())
		}
	}
}

// socketpair returns a non-blocking stream pair with small kernel buffers so
// EAGAIN is reached quickly.
func socketpair(t testing.TB) (w, r int) {
	t.Helper()
	fds, err := unix.Socketpair(unix.AF_UNIX, unix.SOCK_STREAM|unix.SOCK_NONBLOCK|unix.SOCK_CLOEXEC, 0)
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { unix.Close(fds[0]); unix.Close(fds[1]) })
	for _, fd := range fds {
		if err := unix.SetsockoptInt(fd, unix.SOL_SOCKET, unix.SO_SNDBUF, 4096); err != nil {
			t.Fatal(err)
		}
		if err := unix.SetsockoptInt(fd, unix.SOL_SOCKET, unix.SO_RCVBUF, 4096); err != nil {
			t.Fatal(err)
		}
	}
	return fds[0], fds[1]
}

// drain reads everything currently readable from fd into dst.
func drain(t testing.TB, fd int, dst *bytes.Buffer, scratch []byte) {
	t.Helper()
	for {
		n, err := unix.Read(fd, scratch)
		if err == unix.EAGAIN {
			return
		}
		if err != nil {
			t.Fatal(err)
		}
		dst.Write(scratch[:n])
	}
}

// fill queues requests whose bodies carry their sequence number.
func fill(b *Buf, n, body int) []byte {
	for range n {
		p, seq := b.Request(7, 0, body)
		binary.LittleEndian.PutUint32(p, uint32(seq))
	}
	return bytes.Clone(b.Bytes())
}

func TestBufFlushPartialResume(t *testing.T) {
	w, r := socketpair(t)
	b := NewBuf(0)
	want := fill(b, 2000, 60) // 128000 bytes, far beyond the socket buffers
	err := b.flush(w)
	if err != ErrWouldBlock {
		t.Fatal("want ErrWouldBlock, got", err)
	}
	if b.flushed == 0 || b.Len() == 0 || b.Len() >= len(want) {
		t.Fatalf("partial write state: flushed %d pending %d", b.flushed, b.Len())
	}
	if !bytes.Equal(b.Bytes(), want[len(want)-b.Len():]) {
		t.Fatal("Bytes must return the unflushed tail")
	}
	if b.Err() != nil {
		t.Fatal("EAGAIN must not be sticky")
	}
	// Repeated flush on a full socket keeps the state.
	pending := b.Len()
	if err := b.flush(w); err != ErrWouldBlock || b.Len() != pending {
		t.Fatal("second flush on full socket", err, b.Len())
	}
	var got bytes.Buffer
	scratch := make([]byte, 64<<10)
	for rounds := 0; ; rounds++ {
		drain(t, r, &got, scratch)
		err := b.flush(w)
		if err == nil {
			break
		}
		if err != ErrWouldBlock || rounds > 10000 {
			t.Fatal(err, rounds)
		}
	}
	drain(t, r, &got, scratch)
	if !bytes.Equal(got.Bytes(), want) {
		t.Fatalf("stream mismatch: got %d bytes want %d", got.Len(), len(want))
	}
	if b.Len() != 0 || b.flushed != 0 || len(b.b) != 0 || cap(b.b) != 128<<10 {
		t.Fatal("full flush must compact", b.Len(), b.flushed, len(b.b), cap(b.b))
	}
	if err := b.flush(w); err != nil {
		t.Fatal("empty flush", err)
	}
}

func TestBufCompactsAfterPartialFlush(t *testing.T) {
	w, _ := socketpair(t)
	b := NewBuf(0)
	fill(b, 1000, 60) // 64000 bytes in a 64 KiB buffer
	if err := b.flush(w); err != ErrWouldBlock {
		t.Fatal(err)
	}
	flushed, pending := b.flushed, b.Len()
	tail := append([]byte(nil), b.Bytes()...)
	// Does not fit after len(b.b) but fits once flushed bytes are dropped.
	body := (cap(b.b) - pending - 4) &^ 3
	if body+4 <= cap(b.b)-len(b.b) {
		t.Fatal("test setup: request fits without compaction", flushed)
	}
	p, seq := b.Request(9, 0, body)
	if p == nil || cap(b.b) != 64<<10 || b.flushed != 0 {
		t.Fatalf("compaction: seq %d cap %d flushed %d body %d", seq, cap(b.b), b.flushed, body)
	}
	if !bytes.Equal(b.Bytes()[:pending], tail) || b.Bytes()[pending] != 9 {
		t.Fatal("compaction lost pending bytes")
	}
}

func TestBufGrowthKeepsPendingAfterPartialFlush(t *testing.T) {
	w, _ := socketpair(t)
	b := NewBuf(0)
	want := fill(b, 1000, 60)
	if err := b.flush(w); err != ErrWouldBlock {
		t.Fatal(err)
	}
	pending := b.Len()
	b.Request(9, 0, 100000)
	if cap(b.b) != 256<<10 || b.flushed != 0 || !bytes.Equal(b.Bytes()[:pending], want[len(want)-pending:]) {
		t.Fatal("growth with flushed prefix", cap(b.b), b.flushed)
	}
}

func TestBufFlushClosedPeer(t *testing.T) {
	w, r := socketpair(t)
	b := NewBuf(0)
	fill(b, 1, 4)
	unix.Close(r)
	if err := b.flush(w); err != ErrClosed || b.Err() != ErrClosed {
		t.Fatal("want sticky ErrClosed, got", err)
	}
	if p, s := b.Request(1, 0, 0); p != nil || s != 0 {
		t.Fatal("request after dead connection")
	}
	if err := b.flush(w); err != ErrClosed {
		t.Fatal("flush must return sticky error", err)
	}
}

func TestBufFlushOtherError(t *testing.T) {
	b := NewBuf(0)
	fill(b, 1, 4)
	err := b.flush(-1)
	if !errors.Is(err, syscall.EBADF) || b.Err() != err {
		t.Fatal("want sticky EBADF, got", err)
	}
	if b.flush(-1) != err {
		t.Fatal("sticky")
	}
}

func TestBufFlushStickyBeforeWrite(t *testing.T) {
	w, r := socketpair(t)
	b := NewBuf(0)
	fill(b, 1, 4)
	b.Fail(ErrBadLength)
	if err := b.flush(w); err != ErrBadLength {
		t.Fatal(err)
	}
	if n, err := unix.Read(r, make([]byte, 16)); err != unix.EAGAIN {
		t.Fatal("bytes written after sticky error", n)
	}
}

func TestAllocBufFlush(t *testing.T) {
	w, r := socketpair(t)
	b := NewBuf(0)
	scratch := make([]byte, 64<<10)
	n := testing.AllocsPerRun(1000, func() {
		for range 8 {
			b.Request(8, 0, 4)
		}
		if err := b.flush(w); err != nil {
			t.Fatal(err)
		}
		if _, err := unix.Read(r, scratch); err != nil {
			t.Fatal(err)
		}
		b.acked = b.next - 1
	})
	if n != 0 {
		t.Fatal(n)
	}
	// EAGAIN path.
	fill(b, 2000, 60)
	n = testing.AllocsPerRun(100, func() {
		if err := b.flush(w); err != ErrWouldBlock {
			t.Fatal(err)
		}
	})
	if n != 0 {
		t.Fatal("EAGAIN path allocates", n)
	}
}

func BenchmarkBufRequest(b *testing.B) {
	buf := NewBuf(0)
	b.ReportAllocs()
	for b.Loop() {
		buf.Reset()
		buf.acked = buf.next
		buf.Request(1, 0, 4)
	}
}

// BenchmarkBufRequestFlush queues 64 MapWindow-sized requests and flushes
// them to a socketpair per iteration.
func BenchmarkBufRequestFlush(b *testing.B) {
	fds, err := unix.Socketpair(unix.AF_UNIX, unix.SOCK_STREAM|unix.SOCK_NONBLOCK|unix.SOCK_CLOEXEC, 0)
	if err != nil {
		b.Fatal(err)
	}
	defer unix.Close(fds[0])
	defer unix.Close(fds[1])
	buf := NewBuf(0)
	scratch := make([]byte, 64<<10)
	b.ReportAllocs()
	b.SetBytes(64 * 8)
	for b.Loop() {
		for range 64 {
			p, _ := buf.Request(8, 0, 4)
			binary.LittleEndian.PutUint32(p, 0x0300009f)
		}
		if err := buf.flush(fds[0]); err != nil {
			b.Fatal(err)
		}
		if _, err := unix.Read(fds[1], scratch); err != nil {
			b.Fatal(err)
		}
		buf.acked = buf.next - 1
	}
}