channel_test.go

package fastwire

import (
	"testing"
	"time"

	"github.com/marcomoesman/fastwire/internal/rtt"
)

// --- Ack tracking tests ---

func TestAckAdvance(t *testing.T) {
	ch := &channel{mode: Unreliable}
	ch.recvNextDeliver = 1

	// First packet.
	if !ch.recordReceive(1) {
		t.Fatal("seq 1 should be accepted")
	}
	ack, field := ch.ackState()
	if ack != 1 {
		t.Fatalf("recvAck = %d, want 1", ack)
	}
	if field != 0 {
		t.Fatalf("ackField = %032b, want 0", field)
	}

	// Advance to seq 3 (skip 2).
	if !ch.recordReceive(3) {
		t.Fatal("seq 3 should be accepted")
	}
	ack, field = ch.ackState()
	if ack != 3 {
		t.Fatalf("recvAck = %d, want 3", ack)
	}
	// Bit 1 (diff=2) should be set for old recvAck=1, bit 0 (diff=1) unset for missing seq=2.
	// After shift by 2: old ackField=0 << 2 = 0, then set bit (2-1)=bit 1 for old recvAck=1.
	if field&(1<<1) == 0 {
		t.Fatalf("bit for seq 1 should be set, field = %032b", field)
	}
	if field&(1<<0) != 0 {
		t.Fatalf("bit for seq 2 should not be set, field = %032b", field)
	}
}

func TestAckOutOfOrder(t *testing.T) {
	ch := &channel{mode: Unreliable}
	ch.recvNextDeliver = 1

	ch.recordReceive(5)
	ch.recordReceive(3) // out of order
	ch.recordReceive(4) // out of order

	ack, field := ch.ackState()
	if ack != 5 {
		t.Fatalf("recvAck = %d, want 5", ack)
	}
	// seq 3: diff = 5-3 = 2, bit 1
	// seq 4: diff = 5-4 = 1, bit 0
	if field&(1<<1) == 0 {
		t.Fatalf("bit for seq 3 should be set, field = %032b", field)
	}
	if field&(1<<0) == 0 {
		t.Fatalf("bit for seq 4 should be set, field = %032b", field)
	}
}

func TestAckDuplicate(t *testing.T) {
	ch := &channel{mode: Unreliable}
	ch.recvNextDeliver = 1

	if !ch.recordReceive(1) {
		t.Fatal("first receive of seq 1 should be accepted")
	}
	if ch.recordReceive(1) {
		t.Fatal("duplicate seq 1 should be rejected")
	}

	// Also test duplicate of an old packet.
	ch.recordReceive(5)
	if !ch.recordReceive(3) {
		t.Fatal("first receive of seq 3 should be accepted")
	}
	if ch.recordReceive(3) {
		t.Fatal("duplicate seq 3 should be rejected")
	}
}

func TestAckTooOld(t *testing.T) {
	ch := &channel{mode: Unreliable}
	ch.recvNextDeliver = 1

	ch.recordReceive(50)
	// seq 50 - 32 = 18, so seq 17 is too old.
	if ch.recordReceive(17) {
		t.Fatal("seq 17 should be too old (more than 32 behind recvAck 50)")
	}
	// seq 18 is at the boundary (diff=32), should be accepted.
	if !ch.recordReceive(18) {
		t.Fatal("seq 18 should be accepted (exactly 32 behind recvAck 50)")
	}
}

func TestAckBitfieldBoundary(t *testing.T) {
	ch := &channel{mode: Unreliable}
	ch.recvNextDeliver = 1

	// Set recvAck to 33.
	ch.recordReceive(33)

	// seq 1 is diff=32, should be accepted.
	if !ch.recordReceive(1) {
		t.Fatal("seq 1 at boundary should be accepted")
	}

	ack, field := ch.ackState()
	if ack != 33 {
		t.Fatalf("recvAck = %d, want 33", ack)
	}
	// seq 1: diff = 33-1 = 32, bit 31
	if field&(1<<31) == 0 {
		t.Fatalf("bit 31 for seq 1 should be set, field = %032b", field)
	}
}

func TestAckLargeGap(t *testing.T) {
	ch := &channel{mode: Unreliable}
	ch.recvNextDeliver = 1

	ch.recordReceive(1)
	// Large gap — should clear bitfield.
	ch.recordReceive(100)
	_, field := ch.ackState()
	// Only bit for the shifted old recvAck might be set, but diff=99 > 32
	// so the shift clears everything. Actually, when diff > 32, we set ackField=0
	// then we don't set any bit for the old recvAck since it's too far back.
	// Wait, re-reading: diff > 32 → ackField = 0, no bit set for old recvAck.
	// The plan says "shift ackField left by diff, set old recvAck bit".
	// But our implementation handles diff > 32 by clearing the field entirely.
	// That's correct because a 32-bit field can't represent gaps > 32.
	if field != 0 {
		t.Fatalf("bitfield should be 0 after large gap, got %032b", field)
	}
}

func TestAckSeqZeroRejected(t *testing.T) {
	ch := &channel{mode: Unreliable}
	ch.recvNextDeliver = 1
	if ch.recordReceive(0) {
		t.Fatal("sequence 0 should be rejected")
	}
}

// --- Delivery mode tests ---

func TestReliableOrderedReorderAndDrain(t *testing.T) {
	ch := &channel{
		mode:            ReliableOrdered,
		recvBuffer:      make(map[uint32][]byte),
		recvNextDeliver: 1,
	}

	// Receive seq 3, 1, 2 out of order.
	result := ch.deliver(3, []byte("three"))
	if len(result) != 0 {
		t.Fatalf("expected 0 deliveries, got %d", len(result))
	}

	result = ch.deliver(1, []byte("one"))
	if len(result) != 1 || string(result[0]) != "one" {
		t.Fatalf("expected [one], got %v", result)
	}

	result = ch.deliver(2, []byte("two"))
	if len(result) != 2 {
		t.Fatalf("expected 2 deliveries (two, three), got %d", len(result))
	}
	if string(result[0]) != "two" || string(result[1]) != "three" {
		t.Fatalf("expected [two, three], got [%s, %s]", result[0], result[1])
	}
}

func TestReliableUnorderedImmediate(t *testing.T) {
	ch := &channel{mode: ReliableUnordered, recvNextDeliver: 1}

	result := ch.deliver(5, []byte("data"))
	if len(result) != 1 || string(result[0]) != "data" {
		t.Fatalf("expected immediate delivery, got %v", result)
	}
}

func TestUnreliableImmediate(t *testing.T) {
	ch := &channel{mode: Unreliable, recvNextDeliver: 1}

	result := ch.deliver(1, []byte("data"))
	if len(result) != 1 || string(result[0]) != "data" {
		t.Fatalf("expected immediate delivery, got %v", result)
	}
}

func TestUnreliableSequencedDropStale(t *testing.T) {
	ch := &channel{mode: UnreliableSequenced, recvNextDeliver: 1}

	// Receive seq 5.
	result := ch.deliver(5, []byte("five"))
	if len(result) != 1 {
		t.Fatal("expected delivery of seq 5")
	}

	// Receive seq 3 (stale) — should be dropped.
	result = ch.deliver(3, []byte("three"))
	if len(result) != 0 {
		t.Fatal("stale seq 3 should be dropped")
	}

	// Receive seq 5 again (equal, not newer) — should be dropped.
	result = ch.deliver(5, []byte("five-again"))
	if len(result) != 0 {
		t.Fatal("equal seq 5 should be dropped")
	}

	// Receive seq 6 (newer) — should be delivered.
	result = ch.deliver(6, []byte("six"))
	if len(result) != 1 || string(result[0]) != "six" {
		t.Fatalf("expected delivery of seq 6, got %v", result)
	}
}

// --- Process acks tests ---

func TestProcessAcksWithRTT(t *testing.T) {
	ch := &channel{mode: ReliableOrdered, recvBuffer: make(map[uint32][]byte), recvNextDeliver: 1}
	rs := rtt.New()

	now := time.Now()
	ch.addPending(pendingPacket{
		raw:            []byte("packet1"),
		sendTime:       now.Add(-50 * time.Millisecond),
		firstTransmit:  true,
		sequence:       1,
		nextRetransmit: now.Add(1 * time.Second),
	})
	ch.addPending(pendingPacket{
		raw:            []byte("packet2"),
		sendTime:       now.Add(-30 * time.Millisecond),
		firstTransmit:  true,
		sequence:       2,
		nextRetransmit: now.Add(1 * time.Second),
	})

	// Ack seq 2, with bitfield bit 0 set -> also acks seq 1.
	acked := ch.processAcks(2, 0x00000001, rs)

	if len(acked) != 2 {
		t.Fatalf("expected 2 acked, got %d", len(acked))
	}
	if ch.pendingCount() != 0 {
		t.Fatalf("pending count = %d, want 0", ch.pendingCount())
	}

	// RTT should have samples.
	if rs.SRTT() == 0 {
		t.Fatal("SRTT should be > 0 after ack processing")
	}
}

func TestProcessAcksKarnsAlgorithm(t *testing.T) {
	ch := &channel{mode: ReliableOrdered, recvBuffer: make(map[uint32][]byte), recvNextDeliver: 1}
	rs := rtt.New()

	now := time.Now()
	// Retransmitted packet -- firstTransmit = false.
	ch.addPending(pendingPacket{
		raw:            []byte("retransmitted"),
		sendTime:       now.Add(-100 * time.Millisecond),
		firstTransmit:  false,
		sequence:       1,
		nextRetransmit: now.Add(1 * time.Second),
	})

	ch.processAcks(1, 0, rs)

	// RTT should NOT have been updated (Karn's algorithm skips retransmissions).
	if rs.SRTT() != 0 {
		t.Fatalf("SRTT should be 0 (no sample from retransmit), got %v", rs.SRTT())
	}
}

func TestProcessAcksZeroAck(t *testing.T) {
	ch := &channel{mode: ReliableOrdered, recvBuffer: make(map[uint32][]byte), recvNextDeliver: 1}
	rs := rtt.New()

	ch.addPending(pendingPacket{sequence: 1, firstTransmit: true, nextRetransmit: time.Now().Add(time.Second)})

	// ack=0 means "nothing received" -- should not remove anything.
	acked := ch.processAcks(0, 0, rs)
	if len(acked) != 0 {
		t.Fatalf("expected 0 acked with ack=0, got %d", len(acked))
	}
	if ch.pendingCount() != 1 {
		t.Fatalf("pending count = %d, want 1", ch.pendingCount())
	}
}

func TestIsAcked(t *testing.T) {
	tests := []struct {
		name     string
		seq      uint32
		ack      uint32
		ackField uint32
		want     bool
	}{
		{"exact match", 5, 5, 0, true},
		{"ack-1 via bitfield", 4, 5, 0x00000001, true},
		{"ack-1 missing from bitfield", 4, 5, 0x00000000, false},
		{"ack-32 via bitfield", 1, 33, 0x80000000, true},
		{"ack-33 too old", 1, 34, 0xFFFFFFFF, false},
		{"seq > ack", 10, 5, 0xFFFFFFFF, false},
		{"ack is zero", 1, 0, 0xFFFFFFFF, false},
		{"full ackField", 3, 5, 0xFFFFFFFF, true},
		{"seq is zero in bitfield range", 0, 5, 0xFFFFFFFF, true}, // seq 0 never appears in pendingSend; isAcked sees it as ack-5
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			got := isAcked(tt.seq, tt.ack, tt.ackField)
			if got != tt.want {
				t.Errorf("isAcked(%d, %d, 0x%08x) = %v, want %v", tt.seq, tt.ack, tt.ackField, got, tt.want)
			}
		})
	}
}

// --- Retransmission tests ---

func TestRetransmissionScheduling(t *testing.T) {
	ch := &channel{mode: ReliableOrdered, recvBuffer: make(map[uint32][]byte), recvNextDeliver: 1}

	past := time.Now().Add(-100 * time.Millisecond)
	ch.addPending(pendingPacket{
		raw:            []byte("data"),
		sendTime:       past,
		firstTransmit:  true,
		sequence:       1,
		nextRetransmit: past, // already due
	})

	rto := 200 * time.Millisecond
	retransmit, kill := ch.checkRetransmissions(time.Now(), rto, maxRetransmits)
	if kill {
		t.Fatal("should not kill")
	}
	if len(retransmit) != 1 {
		t.Fatalf("expected 1 retransmit, got %d", len(retransmit))
	}
	if retransmit[0].firstTransmit {
		t.Fatal("retransmit should have firstTransmit=false")
	}
	if retransmit[0].retransmitCount != 1 {
		t.Fatalf("retransmitCount = %d, want 1", retransmit[0].retransmitCount)
	}
}

func TestRetransmissionMaxRetriesKill(t *testing.T) {
	ch := &channel{mode: ReliableOrdered, recvBuffer: make(map[uint32][]byte), recvNextDeliver: 1}

	past := time.Now().Add(-100 * time.Millisecond)
	ch.addPending(pendingPacket{
		raw:             []byte("data"),
		sendTime:        past,
		firstTransmit:   false,
		sequence:        1,
		retransmitCount: maxRetransmits, // already at max
		nextRetransmit:  past,
	})

	_, kill := ch.checkRetransmissions(time.Now(), 200*time.Millisecond, maxRetransmits)
	if !kill {
		t.Fatal("should kill when max retransmits exceeded")
	}
}

func TestRetransmissionNotYetDue(t *testing.T) {
	ch := &channel{mode: ReliableOrdered, recvBuffer: make(map[uint32][]byte), recvNextDeliver: 1}

	future := time.Now().Add(10 * time.Second)
	ch.addPending(pendingPacket{
		raw:            []byte("data"),
		sendTime:       time.Now(),
		firstTransmit:  true,
		sequence:       1,
		nextRetransmit: future, // not due yet
	})

	retransmit, kill := ch.checkRetransmissions(time.Now(), 200*time.Millisecond, maxRetransmits)
	if kill {
		t.Fatal("should not kill")
	}
	if len(retransmit) != 0 {
		t.Fatalf("expected 0 retransmits, got %d", len(retransmit))
	}
}

// --- Next sequence tests ---

func TestNextSequenceStartsAtOne(t *testing.T) {
	ch := &channel{mode: Unreliable, recvNextDeliver: 1}
	if seq := ch.nextSequence(); seq != 1 {
		t.Fatalf("first sequence = %d, want 1", seq)
	}
	if seq := ch.nextSequence(); seq != 2 {
		t.Fatalf("second sequence = %d, want 2", seq)
	}
}

// --- Channel layout builder tests ---

func TestDefaultChannelLayout(t *testing.T) {
	layout := DefaultChannelLayout()
	if layout.Len() != 4 {
		t.Fatalf("default layout len = %d, want 4", layout.Len())
	}
	modes := []DeliveryMode{ReliableOrdered, ReliableUnordered, Unreliable, UnreliableSequenced}
	for i, want := range modes {
		if layout.channels[i].Mode != want {
			t.Fatalf("channel %d mode = %d, want %d", i, layout.channels[i].Mode, want)
		}
	}
}

func TestChannelLayoutBuilder(t *testing.T) {
	layout, err := NewChannelLayoutBuilder().
		AddChannel(ReliableOrdered, 0).
		AddChannel(Unreliable, 1).
		Build()
	if err != nil {
		t.Fatal(err)
	}
	if layout.Len() != 2 {
		t.Fatalf("layout len = %d, want 2", layout.Len())
	}
}

func TestChannelLayoutBuilderEmpty(t *testing.T) {
	_, err := NewChannelLayoutBuilder().Build()
	if err != ErrInvalidChannelLayout {
		t.Fatalf("expected ErrInvalidChannelLayout, got %v", err)
	}
}

func TestChannelLayoutBuilderTooMany(t *testing.T) {
	b := NewChannelLayoutBuilder()
	for range 257 {
		b.AddChannel(Unreliable, 0)
	}
	_, err := b.Build()
	if err != ErrInvalidChannelLayout {
		t.Fatalf("expected ErrInvalidChannelLayout, got %v", err)
	}
}

func TestChannelLayoutBuilderMax(t *testing.T) {
	b := NewChannelLayoutBuilder()
	for range 256 {
		b.AddChannel(Unreliable, 0)
	}
	layout, err := b.Build()
	if err != nil {
		t.Fatal(err)
	}
	if layout.Len() != 256 {
		t.Fatalf("layout len = %d, want 256", layout.Len())
	}
}

// --- Pool buffer return tests ---

func TestProcessAcksReturnsPoolBuf(t *testing.T) {
	ch := &channel{mode: ReliableOrdered, recvBuffer: make(map[uint32][]byte), recvNextDeliver: 1}

	// Get a pool buffer and use it as poolBuf.
	buf := getSendBuffer(200)
	now := time.Now()
	ch.addPending(pendingPacket{
		raw:            buf[:100],
		poolBuf:        buf[:cap(buf)],
		sendTime:       now.Add(-50 * time.Millisecond),
		firstTransmit:  true,
		sequence:       1,
		nextRetransmit: now.Add(time.Second),
	})

	acked := ch.processAcks(1, 0, nil)
	if len(acked) != 1 {
		t.Fatalf("expected 1 acked, got %d", len(acked))
	}
	// poolBuf was returned to pool inside processAcks. Verify pool has a buffer.
	buf2 := getSendBuffer(200)
	if cap(buf2) < DefaultMTU {
		t.Fatalf("expected pooled buffer back, got cap %d", cap(buf2))
	}
	putSendBuffer(buf2)
}

func TestProcessAcksNilPoolBuf(t *testing.T) {
	ch := &channel{mode: ReliableOrdered, recvBuffer: make(map[uint32][]byte), recvNextDeliver: 1}

	now := time.Now()
	ch.addPending(pendingPacket{
		raw:            []byte("data"),
		poolBuf:        nil, // not pooled
		sendTime:       now.Add(-50 * time.Millisecond),
		firstTransmit:  true,
		sequence:       1,
		nextRetransmit: now.Add(time.Second),
	})

	// Should not panic with nil poolBuf.
	acked := ch.processAcks(1, 0, nil)
	if len(acked) != 1 {
		t.Fatalf("expected 1 acked, got %d", len(acked))
	}
}

func TestReleasePendingBuffers(t *testing.T) {
	ch := &channel{mode: ReliableOrdered, recvBuffer: make(map[uint32][]byte), recvNextDeliver: 1}

	// Add several pending packets with pool buffers.
	for i := range 5 {
		buf := getSendBuffer(200)
		ch.addPending(pendingPacket{
			raw:            buf[:100],
			poolBuf:        buf[:cap(buf)],
			sequence:       uint32(i + 1),
			nextRetransmit: time.Now().Add(time.Second),
		})
	}

	if ch.pendingCount() != 5 {
		t.Fatalf("pending count = %d, want 5", ch.pendingCount())
	}

	ch.releasePendingBuffers()

	if ch.pendingCount() != 0 {
		t.Fatalf("pending count after release = %d, want 0", ch.pendingCount())
	}
}

func TestReleasePendingBuffersEmpty(t *testing.T) {
	ch := &channel{mode: Unreliable, recvNextDeliver: 1}
	// Should not panic on empty channel.
	ch.releasePendingBuffers()
}

func TestReleasePendingBuffersClearsRecvBuffer(t *testing.T) {
	ch := &channel{mode: ReliableOrdered, recvBuffer: make(map[uint32][]byte), recvNextDeliver: 1}

	// Simulate out-of-order buffered payloads.
	ch.recvBuffer[5] = []byte("payload-5")
	ch.recvBuffer[6] = []byte("payload-6")
	ch.recvBuffer[8] = []byte("payload-8")

	if len(ch.recvBuffer) != 3 {
		t.Fatalf("recvBuffer len = %d, want 3", len(ch.recvBuffer))
	}

	ch.releasePendingBuffers()

	if ch.recvBuffer != nil {
		t.Fatalf("recvBuffer should be nil after release, got len %d", len(ch.recvBuffer))
	}
}

// --- newChannels test ---

func TestNewChannels(t *testing.T) {
	layout := DefaultChannelLayout()
	chs := newChannels(layout, DefaultMaxReorderWindow)
	if len(chs) != 4 {
		t.Fatalf("len(channels) = %d, want 4", len(chs))
	}

	// ReliableOrdered should have a recvBuffer.
	if chs[0].recvBuffer == nil {
		t.Fatal("channel 0 (ReliableOrdered) should have recvBuffer")
	}
	// Others should not.
	for i := 1; i < 4; i++ {
		if chs[i].recvBuffer != nil {
			t.Fatalf("channel %d should not have recvBuffer", i)
		}
	}
}

// --- Reorder window bound ---

func TestReliableOrderedReorderWindowCap(t *testing.T) {
	const window = 8
	chs := newChannels(DefaultChannelLayout(), window)
	ch := chs[0] // ReliableOrdered

	// Seq 1 fills recvNextDeliver slot cleanly.
	if !ch.recordReceive(1) {
		t.Fatal("seq 1 should be accepted")
	}
	ch.deliver(1, []byte("x"))

	// Now next-expected is 2. Accept up to window-1 ahead (seq ≤ 9).
	if !ch.recordReceive(9) {
		t.Fatal("seq 9 (window edge) should be accepted")
	}
	// Seq 10 is exactly at the window boundary — must be dropped.
	if ch.recordReceive(10) {
		t.Fatal("seq 10 (beyond window) should be rejected")
	}
	// Far-future sequence also dropped.
	if ch.recordReceive(10_000) {
		t.Fatal("far-future seq should be rejected")
	}
}

func TestReorderWindowDisabledWhenZero(t *testing.T) {
	chs := newChannels(DefaultChannelLayout(), 0)
	ch := chs[0]
	// With no cap, any seq should land.
	if !ch.recordReceive(1_000_000) {
		t.Fatal("seq 1M should be accepted when window is unbounded")
	}
}