package xturbo

import (
	"bytes"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"testing"
	"time"

	"bnema.dev/xturbo/internal/xtest"
	"golang.org/x/sys/unix"
)

// dispatchConn returns a Conn over one end of a non-blocking socketpair with
// default kernel buffers, and the other end acting as the server. The XFIXES
// and Generic Event extensions are negotiated at the golden capture's bases.
func dispatchConn(t testing.TB) (*Conn, 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)
	}
	c := newConnState(fds[0], nil, setupInfo{})
	t.Cleanup(func() { _ = c.Close(); _ = unix.Close(fds[1]) })
	c.ext.Set(ExtXFixes, 138, 87, 140)
	c.ext.Set(ExtGenericEvent, 0x83, 0, 0)
	c.buf.next = 1 << 40 // every widened seq in the tests is acknowledgeable
	return c, fds[1]
}

// Packet builders. seq is the 16-bit wire value.
func pktEvent(code uint8, seq uint16, fill byte) []byte {
	p := bytes.Repeat([]byte{fill}, 32)
	p[0] = code
	binary.LittleEndian.PutUint16(p[2:], seq)
	return p
}

func pktReply(seq uint16, data int) []byte {
	words := (data + 3) / 4
	p := make([]byte, 32+4*words)
	p[0], p[1] = 1, 7
	binary.LittleEndian.PutUint16(p[2:], seq)
	binary.LittleEndian.PutUint32(p[4:], uint32(words))
	for i := 8; i < len(p); i++ {
		p[i] = byte(i * 7)
	}
	return p
}

func pktError(code uint8, seq uint16, major uint8) []byte {
	p := make([]byte, 32)
	p[1] = code
	binary.LittleEndian.PutUint16(p[2:], seq)
	binary.LittleEndian.PutUint32(p[4:], 0xdead)
	p[10] = major
	return p
}

func pktXGE(seq uint16, evtype uint16, tail int) []byte {
	words := (tail + 3) / 4
	p := make([]byte, 32+4*words)
	p[0], p[1] = codeGenericEvent, 0x83
	binary.LittleEndian.PutUint16(p[2:], seq)
	binary.LittleEndian.PutUint32(p[4:], uint32(words))
	binary.LittleEndian.PutUint16(p[8:], evtype)
	for i := 32; i < len(p); i++ {
		p[i] = byte(i)
	}
	return p
}

// record is one callback as seen by a test handler; data is a copy of the
// packet bytes (Raw, plus Tail for events).
type record struct {
	class byte // 'e', 'r', 'x'
	seq   Seq
	kind  uint32
	data  []byte
}

func (r record) String() string {
	return fmt.Sprintf("%c seq %d kind %#x %d bytes", r.class, r.seq, r.kind, len(r.data))
}

type recorder struct {
	got []record
	h   Handler
}

// newRecorder records every callback of c. Replies reach OnReply only when
// registered: on each MapRequest event (the first packet of scriptedStream)
// it registers the stream's replies, at ev.Seq+1, +4 and +6.
func newRecorder(c *Conn) *recorder {
	r := &recorder{}
	r.h = Handler{
		OnEvent: func(ev *Event) {
			expectScripted(c, ev)
			r.got = append(r.got, record{'e', ev.Seq, uint32(ev.Kind), append(bytes.Clone(ev.Raw[:]), ev.Tail...)})
		},
		OnReply: func(rp *Reply) {
			r.got = append(r.got, record{'r', rp.Seq, uint32(rp.Kind), bytes.Clone(rp.Raw)})
		},
		OnError: func(e *XError) {
			r.got = append(r.got, record{'x', e.Seq, uint32(e.Code), bytes.Clone(e.Raw[:])})
		},
	}
	return r
}

// feed writes data to the server end, dispatching whenever the socket fills,
// then dispatches what remains.
func feed(t testing.TB, c *Conn, srv int, data []byte, h *Handler) {
	t.Helper()
	for len(data) > 0 {
		n, err := unix.Write(srv, data)
		if err == unix.EAGAIN {
			n = 0
		} else if err != nil {
			t.Fatal(err)
		}
		data = data[n:]
		pump(t, c, h)
	}
	pump(t, c, h)
}

// pump dispatches until the socket holds nothing more: one Dispatch reads a
// bounded amount.
func pump(t testing.TB, c *Conn, h *Handler) {
	t.Helper()
	for {
		if err := c.Dispatch(h); err != nil {
			t.Fatal(err)
		}
		fds := []unix.PollFd{{Fd: int32(c.fd), Events: unix.POLLIN}}
		if n, _ := unix.Poll(fds, 0); n == 0 {
			return
		}
	}
}

func golden(t testing.TB, name string) []byte {
	t.Helper()
	raw, err := xtest.LoadHex("testdata/golden/" + name + ".hex")
	if err != nil {
		t.Fatal(err)
	}
	return raw
}

// scriptedStream is a mixed stream of goldens and synthetic packets with
// increasing wire sequence numbers, and the callbacks it must produce.
func scriptedStream(t testing.TB) ([]byte, []record) {
	var stream []byte
	var want []record
	add := func(class byte, kind uint32, p []byte, seq uint16) {
		binary.LittleEndian.PutUint16(p[2:], seq)
		stream = append(stream, p...)
		want = append(want, record{class, Seq(seq), kind, p})
	}
	add('e', 20, golden(t, "xproto/MapRequest.event"), 3)
	add('r', uint32(kindUser), golden(t, "xproto/GetGeometry.reply"), 4)
	add('x', 3, golden(t, "xproto/BadWindow.error"), 5)
	add('e', uint32(ExtGenericEvent)<<8|0x81, golden(t, "xproto/GeGeneric.event"), 6)
	add('e', uint32(ExtXFixes)<<8, golden(t, "xfixes/SelectionNotify.event"), 6)
	add('r', uint32(kindUser), golden(t, "xproto/GetProperty.reply"), 7)
	km := golden(t, "xproto/KeymapNotify.event") // no sequence: inherits 7
	stream = append(stream, km...)
	want = append(want, record{'e', 7, codeKeymapNotify, km})
	add('e', uint32(KindUnknown), pktEvent(80, 0, 0x11), 8) // unbound code
	add('e', uint32(ExtGenericEvent)<<8|0x80, pktXGE(0, 0, 0), 8)
	add('r', uint32(kindUser), pktReply(0, 0), 9)
	sent := pktEvent(0x80|33, 0, 0x22) // SendEvent ClientMessage
	add('e', 33, sent, 9)
	return stream, want
}

// expectScripted registers the replies of scriptedStream when its first
// event arrives. Allocation-free, so the alloc tests use it too.
func expectScripted(c *Conn, ev *Event) {
	if ev.Kind != 20 {
		return
	}
	for _, d := range [...]Seq{1, 4, 6} {
		if err := c.Expect(ev.Seq+d, kindUser, uint32(d), 0, 0, 0); err != nil {
			panic(err)
		}
	}
}

// expect registers seqs with kind kindUser.
func expect(t testing.TB, c *Conn, seqs ...Seq) {
	t.Helper()
	for _, s := range seqs {
		if err := c.Expect(s, kindUser, 0, 0, 0, 0); err != nil {
			t.Fatal(err)
		}
	}
}

func checkRecords(t testing.TB, got, want []record) {
	t.Helper()
	if len(got) != len(want) {
		t.Fatalf("%d callbacks, want %d: %v", len(got), len(want), got)
	}
	for i := range want {
		g, w := got[i], want[i]
		if g.class != w.class || g.seq != w.seq || g.kind != w.kind || !bytes.Equal(g.data, w.data) {
			t.Fatalf("callback %d: %v, want %v", i, g, w)
		}
	}
}

func TestDispatchStream(t *testing.T) {
	stream, want := scriptedStream(t)
	t.Run("one read", func(t *testing.T) {
		c, srv := dispatchConn(t)
		r := newRecorder(c)
		feed(t, c, srv, stream, &r.h)
		checkRecords(t, r.got, want)
		if c.buf.acked != 9 || c.seq != 9 {
			t.Fatalf("acked %d seq %d", c.buf.acked, c.seq)
		}
		if st := c.Stats(); st.UnknownEvents != 1 || st.DroppedReplies != 0 || st.UnhandledErrors != 0 || st.MaxRingUse != 3 {
			t.Fatalf("%+v", st)
		}
	})
	t.Run("split at every byte", func(t *testing.T) {
		for k := range len(stream) + 1 {
			c, srv := dispatchConn(t)
			r := newRecorder(c)
			feed(t, c, srv, stream[:k], &r.h)
			feed(t, c, srv, stream[k:], &r.h)
			checkRecords(t, r.got, want)
		}
	})
	t.Run("byte by byte", func(t *testing.T) {
		c, srv := dispatchConn(t)
		r := newRecorder(c)
		for i := range stream {
			feed(t, c, srv, stream[i:i+1], &r.h)
		}
		checkRecords(t, r.got, want)
	})
	t.Run("nil handler and nil fields", func(t *testing.T) {
		c, srv := dispatchConn(t)
		feed(t, c, srv, stream, nil)
		feed(t, c, srv, stream, &Handler{})
		if c.rpos != c.rend {
			t.Fatal("stream not consumed")
		}
	})
}

func TestDispatchSeqWidening(t *testing.T) {
	c, srv := dispatchConn(t)
	r := newRecorder(c)
	// Wire seqs climbing through two 16-bit wraps.
	wires := []uint16{1, 0x7fff, 0xfffe, 0xffff, 0, 3, 0x8000, 0xffff, 2}
	wantSeq := []Seq{1, 0x7fff, 0xfffe, 0xffff, 0x10000, 0x10003, 0x18000, 0x1ffff, 0x20002}
	var stream []byte
	for i, w := range wires {
		switch i % 3 {
		case 0:
			stream = append(stream, pktReply(w, 4)...)
		case 1:
			stream = append(stream, pktEvent(22, w, 0)...)
		default:
			stream = append(stream, pktError(3, w, 12)...)
		}
	}
	stream = append(stream, golden(t, "xproto/KeymapNotify.event")...)
	expect(t, c, 1, 0xffff, 0x18000)
	feed(t, c, srv, stream, &r.h)
	if len(r.got) != len(wires)+1 {
		t.Fatal(r.got)
	}
	for i, s := range wantSeq {
		if r.got[i].seq != s {
			t.Fatalf("packet %d: seq %#x, want %#x", i, r.got[i].seq, s)
		}
	}
	if last := r.got[len(wires)]; last.seq != 0x20002 || c.buf.acked != 0x20002 {
		t.Fatalf("KeymapNotify seq %#x, acked %#x", last.seq, c.buf.acked)
	}

	// ack never passes requests that were not sent.
	c2, srv2 := dispatchConn(t)
	c2.buf.next = 5
	if err := c2.expect(9, kindUser, [4]uint32{}); err != nil { // internal: Expect rejects unsent seqs
		t.Fatal(err)
	}
	feed(t, c2, srv2, pktReply(9, 0), nil)
	if c2.buf.acked != 0 || c2.seq != 9 {
		t.Fatalf("acked %d seq %d", c2.buf.acked, c2.seq)
	}
}

func TestDispatchGrowth(t *testing.T) {
	c, srv := dispatchConn(t)
	r := newRecorder(c)
	big := pktReply(1, 300<<10) // larger than the 64 KiB ring
	xge := pktXGE(2, 5, 100<<10)
	stream := append(append(append([]byte{}, pktEvent(22, 1, 0)...), big...), xge...)
	expect(t, c, 1)
	feed(t, c, srv, stream, &r.h)
	if len(r.got) != 3 || !bytes.Equal(r.got[1].data, big) || !bytes.Equal(r.got[2].data, xge) {
		t.Fatalf("%v", r.got)
	}
	if len(c.rb) != 512<<10 {
		t.Fatalf("ring %d bytes, want 512 KiB", len(c.rb))
	}
}

func TestDispatchMaxPacket(t *testing.T) {
	t.Run("exactly MaxPacket", func(t *testing.T) {
		c, srv := dispatchConn(t)
		c.maxPacket = 256 << 10
		r := newRecorder(c)
		p := pktReply(1, c.maxPacket-32)
		expect(t, c, 1)
		feed(t, c, srv, p, &r.h)
		if len(r.got) != 1 || len(r.got[0].data) != c.maxPacket || len(c.rb) != c.maxPacket {
			t.Fatalf("%v ring %d", r.got, len(c.rb))
		}
	})
	for _, tc := range []struct {
		name string
		p    []byte
	}{
		{"reply", pktReply(1, 0)},
		{"XGE", pktXGE(1, 0, 0)},
	} {
		t.Run(tc.name+" beyond MaxPacket", func(t *testing.T) {
			c, srv := dispatchConn(t)
			r := newRecorder(c)
			p := tc.p
			binary.LittleEndian.PutUint32(p[4:], (MaxPacket-32)/4+1)
			stream := append(pktEvent(22, 1, 0), p...)
			if _, err := unix.Write(srv, stream); err != nil {
				t.Fatal(err)
			}
			if err := c.Dispatch(&r.h); !errors.Is(err, ErrPacketTooLarge) {
				t.Fatal(err)
			}
			if len(r.got) != 1 || !errors.Is(c.Dispatch(&r.h), ErrPacketTooLarge) || !errors.Is(c.buf.Err(), ErrPacketTooLarge) {
				t.Fatal("not sticky")
			}
		})
	}
	t.Run("length overflowing 32 bits", func(t *testing.T) {
		if n, ok := packetSize(bytes.Repeat([]byte{1, 0, 0, 0, 0xff, 0xff, 0xff, 0xff}, 4), MaxPacket); ok {
			t.Fatal(n)
		}
	})
}

func TestDispatchEOF(t *testing.T) {
	for _, tail := range [][]byte{nil, pktReply(2, 64)[:40]} {
		c, srv := dispatchConn(t)
		r := newRecorder(c)
		data := append(pktEvent(22, 1, 0), tail...)
		if _, err := unix.Write(srv, data); err != nil {
			t.Fatal(err)
		}
		unix.Shutdown(srv, unix.SHUT_WR)
		if err := c.Dispatch(&r.h); err != io.EOF {
			t.Fatal(err)
		}
		if len(r.got) != 1 || c.Dispatch(&r.h) != io.EOF || !errors.Is(c.buf.Err(), ErrClosed) {
			t.Fatal("EOF not sticky")
		}
	}
	c, _ := dispatchConn(t)
	c.Close()
	if !errors.Is(c.Dispatch(nil), ErrClosed) {
		t.Fatal("Dispatch after Close")
	}
}

func TestDispatchCloseInCallback(t *testing.T) {
	c, srv := dispatchConn(t)
	n := 0
	h := Handler{OnEvent: func(*Event) { n++; c.Close() }}
	if _, err := unix.Write(srv, append(pktEvent(22, 1, 0), pktEvent(22, 2, 0)...)); err != nil {
		t.Fatal(err)
	}
	if err := c.Dispatch(&h); !errors.Is(err, ErrClosed) || n != 1 {
		t.Fatal(err, n)
	}
}

func TestDispatchReentrant(t *testing.T) {
	c, srv := dispatchConn(t)
	var nested error
	var got []Seq
	h := Handler{OnEvent: func(ev *Event) {
		got = append(got, ev.Seq)
		if len(got) == 1 {
			nested = c.Dispatch(nil)
		}
	}}
	var stream []byte
	for i := range 5 {
		stream = append(stream, pktEvent(22, uint16(i+1), 0)...)
	}
	if _, err := unix.Write(srv, stream); err != nil {
		t.Fatal(err)
	}
	if err := c.Dispatch(&h); err != nil {
		t.Fatal(err)
	}
	if !errors.Is(nested, errReentrantDispatch) || len(got) != 5 || got[4] != 5 {
		t.Fatal(nested, got)
	}
	if err := c.Dispatch(&h); err != nil || c.dispatching { // usable afterwards
		t.Fatal(err)
	}
}

func TestDispatchBounded(t *testing.T) {
	c, srv := dispatchConn(t)
	stop := make(chan struct{})
	done := make(chan struct{})
	go func() { // keep the socket full until told to stop
		defer close(done)
		buf := bytes.Repeat(pktEvent(22, 1, 0), 2048)
		for {
			select {
			case <-stop:
				return
			default:
			}
			if _, err := unix.Write(srv, buf); err != nil && err != unix.EAGAIN {
				return
			}
		}
	}()
	defer func() { close(stop); <-done }()
	events := 0
	h := Handler{OnEvent: func(*Event) { events++ }}
	fds := []unix.PollFd{{Fd: int32(c.fd), Events: unix.POLLIN}}
	for n, _ := unix.Poll(fds, 1000); n == 0; n, _ = unix.Poll(fds, 1000) {
		t.Fatal("writer never filled the socket")
	}
	res := make(chan error, 1)
	go func() { res <- c.Dispatch(&h) }()
	select {
	case err := <-res:
		if err != nil {
			t.Fatal(err)
		}
	case <-time.After(5 * time.Second):
		t.Fatal("Dispatch did not return under a flooding server")
	}
	// one budget of reads, plus at most the bytes already in the ring
	if max := 2 * len(c.rb) / 32; events == 0 || events > max {
		t.Fatalf("%d events, want 1..%d", events, max)
	}
}

func TestDispatchReservedKindError(t *testing.T) {
	c, srv := dispatchConn(t)
	r := newRecorder(c)
	if err := c.expect(1, kindDiscard, [4]uint32{}); err != nil {
		t.Fatal(err)
	}
	feed(t, c, srv, pktError(3, 1, 0), &r.h)
	if len(r.got) != 0 || c.Stats().UnhandledErrors != 0 {
		t.Fatalf("reserved-kind error reached OnError: %v", r.got)
	}
	if c.live != 0 {
		t.Fatal("slot not freed")
	}
}

func TestDispatchClearsBorrowedViews(t *testing.T) {
	c, srv := dispatchConn(t)
	var h Handler
	h.OnReply = func(*Reply) {}
	expect(t, c, 2)
	stream := append(append(pktReply(1, 64), pktReply(2, 64)...), pktXGE(3, 5, 64)...)
	feed(t, c, srv, stream, &h) // reply 1 is dropped, 2 delivered, 3 an event with Tail
	if c.reply.Raw != nil || c.ev.Tail != nil {
		t.Fatalf("retained views: reply %d ev %d", len(c.reply.Raw), len(c.ev.Tail))
	}
	feed(t, c, srv, pktReply(4, 8), nil) // dropped reply
	if c.reply.Raw != nil {
		t.Fatal("dropped reply retained")
	}
}

// recordedStream repeats the scripted stream (events, replies, errors) up to
// about size bytes.
func recordedStream(t testing.TB, size int) ([]byte, int) {
	unit, want := scriptedStream(t)
	var s []byte
	n := 0
	for len(s) < size {
		s = append(s, unit...)
		n += len(want)
	}
	return s, n
}

func TestAllocDispatch(t *testing.T) {
	stream, packets := recordedStream(t, 1<<20)
	c, srv := dispatchConn(t)
	var events, replies, errs int
	var sink byte
	h := Handler{
		OnEvent: func(ev *Event) { expectScripted(c, ev); events++; sink ^= ev.Raw[1]; sink ^= byte(len(ev.Tail)) },
		OnReply: func(r *Reply) { replies++; sink ^= r.Raw[len(r.Raw)-1] },
		OnError: func(e *XError) { errs++; sink ^= e.Code },
	}
	feed(t, c, srv, stream, &h) // warm the ring
	if n := testing.AllocsPerRun(5, func() { feed(t, c, srv, stream, &h) }); n != 0 {
		t.Fatalf("Dispatch allocates %v times per 1 MiB", n)
	}
	if events+replies+errs != 7*packets || errs == 0 || replies == 0 || c.Stats().DroppedReplies != 0 {
		t.Fatal(events, replies, errs, c.Stats())
	}
	_ = sink
}

// mixedPackets is a 10k-packet stream: core events, replies, errors, XFIXES
// and XGE events.
func mixedPackets(t testing.TB) []byte {
	var s []byte
	for i := range 10000 {
		seq := uint16(i)
		switch i % 5 {
		case 0:
			s = append(s, pktEvent(22, seq, 0)...)
		case 1:
			s = append(s, pktReply(seq, 24)...)
		case 2:
			s = append(s, pktError(3, seq, 12)...)
		case 3:
			s = append(s, pktEvent(87, seq, 1)...)
		default:
			s = append(s, pktXGE(seq, 1, 8)...)
		}
	}
	return s
}

func BenchmarkDispatch10k(b *testing.B) {
	stream := mixedPackets(b)
	c, srv := dispatchConn(b)
	var sink int
	h := Handler{
		OnEvent: func(ev *Event) {
			if ev.Kind == 22 {
				_ = c.Expect(ev.Seq+1, kindUser, 0, 0, 0, 0)
			}
			sink += int(ev.Kind)
		},
		OnReply: func(r *Reply) { sink += len(r.Raw) },
		OnError: func(e *XError) { sink += int(e.Code) },
	}
	feed(b, c, srv, stream, &h)
	b.SetBytes(int64(len(stream)))
	b.ReportAllocs()
	for b.Loop() {
		feed(b, c, srv, stream, &h)
	}
	if c.Stats().DroppedReplies != 0 {
		b.Fatal(c.Stats())
	}
	_ = sink
}

// FuzzDispatchStream feeds an arbitrary byte stream in arbitrary chunks.
// Invariants: no panic; the callbacks' bytes concatenate to the stream's
// complete packets minus the (unregistered, so dropped) replies, which
// Stats.DroppedReplies counts; Seq never decreases; Dispatch ends with io.EOF
// (or ErrPacketTooLarge) after the server hangs up.
func FuzzDispatchStream(f *testing.F) {
	stream, _ := scriptedStream(f)
	f.Add(stream, uint16(7))
	f.Add(stream, uint16(0x2000|7))
	f.Add(stream, uint16(0))
	f.Add(append(pktReply(1, 1024), pktXGE(2, 1, 64)...), uint16(33))
	f.Fuzz(func(t *testing.T, stream []byte, chunk uint16) {
		if len(stream) > 1<<16 {
			return
		}
		c, srv := dispatchConn(t)
		c.maxPacket = 8 << 10
		c.rb = make([]byte, packetHeadSize+int(chunk>>8)%96) // tiny ring: fuzz growth
		var got []byte
		var last Seq
		mono := func(s Seq) {
			if s < last {
				t.Fatalf("seq %d after %d", s, last)
			}
			last = s
		}
		h := Handler{
			OnEvent: func(ev *Event) { mono(ev.Seq); got = append(append(got, ev.Raw[:]...), ev.Tail...) },
			OnReply: func(r *Reply) { mono(r.Seq); got = append(got, r.Raw...) },
			OnError: func(e *XError) { mono(e.Seq); got = append(got, e.Raw[:]...) },
		}
		step := int(chunk)%97 + 1
		var err error
		for p := stream; len(p) > 0 && err == nil; {
			k := min(step, len(p))
			if _, werr := unix.Write(srv, p[:k]); werr != nil {
				t.Fatal(werr)
			}
			p = p[k:]
			err = c.Dispatch(&h)
		}
		unix.Shutdown(srv, unix.SHUT_WR)
		for err == nil {
			err = c.Dispatch(&h)
		}
		if err != io.EOF && !errors.Is(err, ErrPacketTooLarge) {
			t.Fatal(err)
		}
		var want []byte
		var replies uint64
		for p := stream; len(p) >= 32; {
			n, ok := packetSize(p, c.maxPacket)
			if !ok || n > len(p) {
				break
			}
			if p[0] == 1 {
				replies++
			} else {
				want = append(want, p[:n]...)
			}
			p = p[n:]
		}
		if !bytes.Equal(got, want) || c.Stats().DroppedReplies != replies {
			t.Fatalf("callbacks diverge from the stream: %d/%d bytes, %d dropped of %d replies", len(got), len(want), c.Stats().DroppedReplies, replies)
		}
	})
}