package relay

import (
	"net"
	"sync"
	"testing"
	"time"
)

func TestRelayConnGatesEIPPorts(t *testing.T) {
	inner := &fakePacketConn{}
	ports := &portSet{m: map[int]bool{50000: true}}
	eip := net.ParseIP("203.0.113.5")
	now := time.Unix(1_800_000_000, 0)
	clock := func() time.Time { return now }
	var in int
	c := newRelayConn(inner, &Binding{ID: Identity{"abc", "meet"}}, eip, ports, NewBucket(1<<40, clock), NewBucket(1<<40, clock),
		func(_ Identity, n int) { in += n }, nil, func(*relayConn) {})
	ok := &net.UDPAddr{IP: eip, Port: 50000}
	bad := &net.UDPAddr{IP: eip, Port: 3478}
	other := &net.UDPAddr{IP: net.ParseIP("198.51.100.9"), Port: 4444}
	c.WriteTo([]byte("abc"), ok)
	c.WriteTo([]byte("abc"), bad)
	c.WriteTo([]byte("abcd"), other)
	if len(inner.written) != 2 || in != 7 {
		t.Fatalf("written=%d in=%d, want 2 packets and 7 bytes", len(inner.written), in)
	}
}

func TestRelayConnBucketDrops(t *testing.T) {
	inner := &fakePacketConn{}
	now := time.Unix(1_800_000_000, 0)
	clock := func() time.Time { return now }
	c := newRelayConn(inner, &Binding{}, nil, &portSet{m: map[int]bool{}}, NewBucket(80, clock), NewBucket(1<<40, clock), func(Identity, int) {}, nil, func(*relayConn) {})
	dst := &net.UDPAddr{IP: net.ParseIP("198.51.100.9"), Port: 4444}
	c.WriteTo(make([]byte, 10), dst) // 10 bytes = the whole second
	c.WriteTo(make([]byte, 1), dst)
	if len(inner.written) != 1 {
		t.Fatalf("written = %d, want 1 (second dropped by bucket)", len(inner.written))
	}
}

func TestRelayConnBoxBucketRefusalDoesNotDebitApp(t *testing.T) {
	now := time.Unix(1_800_000_000, 0)
	clock := func() time.Time { return now }
	appBucket := NewBucket(88, clock) // 11 bytes/s: exactly 1 byte of headroom after a 10-byte write
	boxBucket := NewBucket(80, clock) // 10 bytes/s
	inner := &fakePacketConn{}
	c := newRelayConn(inner, &Binding{}, nil, &portSet{m: map[int]bool{}}, appBucket, boxBucket, func(Identity, int) {}, nil, func(*relayConn) {})
	dst := &net.UDPAddr{IP: net.ParseIP("198.51.100.9"), Port: 4444}

	c.WriteTo(make([]byte, 10), dst) // spends the whole second on both buckets: app 11->1, box 10->0
	if len(inner.written) != 1 {
		t.Fatalf("written = %d, want 1 after the first write", len(inner.written))
	}

	c.WriteTo(make([]byte, 1), dst) // box has 0 left: refused. With the bug, appBucket.Allow already spent app's last token here.
	if len(inner.written) != 1 {
		t.Fatalf("written = %d, want 1 (dropped by the box bucket)", len(inner.written))
	}

	// A fresh, effectively unlimited box bucket, but the SAME app bucket.
	// If the earlier box refusal wrongly debited app (the bug), app is now
	// at 0 and this write is refused too. Fixed, app still has its 1 token.
	freshBox := NewBucket(1<<40, clock)
	c2 := newRelayConn(inner, &Binding{}, nil, &portSet{m: map[int]bool{}}, appBucket, freshBox, func(Identity, int) {}, nil, func(*relayConn) {})
	c2.WriteTo(make([]byte, 1), dst)
	if len(inner.written) != 2 {
		t.Fatalf("written = %d, want 2 (app bucket still had its last token after the box refusal)", len(inner.written))
	}
}

func TestRelayConnCloseOnce(t *testing.T) {
	n := 0
	c := newRelayConn(&fakePacketConn{}, &Binding{}, nil, &portSet{m: map[int]bool{}}, nil, nil, nil, nil, func(*relayConn) { n++ })
	c.Close()
	c.Close()
	if n != 1 {
		t.Fatalf("onClose ran %d times", n)
	}
}

func TestTypedSrc(t *testing.T) {
	u := &net.UDPAddr{IP: net.ParseIP("192.0.2.1"), Port: 5000}
	tc := &net.TCPAddr{IP: net.ParseIP("192.0.2.1"), Port: 5000}
	if typedSrc(u) == typedSrc(tc) {
		t.Fatal("udp and tcp with the same address must not collide")
	}
}

// fakePacketConn records writes; ReadFrom blocks until Close.
type fakePacketConn struct {
	mu      sync.Mutex
	written [][]byte
	closed  chan struct{}
	once    sync.Once
}

func (f *fakePacketConn) init() { f.once.Do(func() { f.closed = make(chan struct{}) }) }
func (f *fakePacketConn) ReadFrom(b []byte) (int, net.Addr, error) {
	f.init()
	<-f.closed
	return 0, nil, net.ErrClosed
}
func (f *fakePacketConn) WriteTo(b []byte, a net.Addr) (int, error) {
	f.mu.Lock()
	defer f.mu.Unlock()
	f.written = append(f.written, append([]byte(nil), b...))
	return len(b), nil
}
func (f *fakePacketConn) Close() error {
	f.init()
	select {
	case <-f.closed:
	default:
		close(f.closed)
	}
	return nil
}
func (f *fakePacketConn) LocalAddr() net.Addr              { return &net.UDPAddr{IP: net.IPv4zero, Port: 50000} }
func (f *fakePacketConn) SetDeadline(time.Time) error      { return nil }
func (f *fakePacketConn) SetReadDeadline(time.Time) error  { return nil }
func (f *fakePacketConn) SetWriteDeadline(time.Time) error { return nil }