integration_test.go

package netpipe

import (
	"bytes"
	"fmt"
	"sync"
	"testing"
	"time"
)

// findFreePort returns an unused port by letting the OS assign one.
func findFreePort() int {
	// use a port in the high range to avoid conflicts
	return 0 // will use specific ports per test
}

// ---------------------------------------------------------------------------
// TCP server ↔ client round-trip
// ---------------------------------------------------------------------------

func TestIntegrationTCPRoundTrip(t *testing.T) {
	port := 17100

	server := NewServer(Config{Port: port, ConnectTimeout: 5 * time.Second})

	var received []byte
	var mu sync.Mutex
	done := make(chan struct{})

	server.OnData(func(id string, data []byte) {
		mu.Lock()
		received = append(received, data...)
		mu.Unlock()
		// echo back
		server.SendTo(id, data)
	})

	go server.Listen()
	defer server.Close()
	time.Sleep(50 * time.Millisecond) // let server start

	client := NewClient(Config{Port: port})

	var clientReceived []byte
	client.OnData(func(data []byte) {
		mu.Lock()
		clientReceived = append(clientReceived, data...)
		mu.Unlock()
		close(done)
	})

	if err := client.Connect(); err != nil {
		t.Fatalf("connect: %v", err)
	}
	defer client.Disconnect()

	go client.Listen()
	time.Sleep(50 * time.Millisecond) // let session frame arrive

	msg := []byte("hello from integration test")
	client.Send(msg)

	select {
	case <-done:
	case <-time.After(3 * time.Second):
		t.Fatal("timeout waiting for echo")
	}

	mu.Lock()
	defer mu.Unlock()
	if !bytes.Equal(clientReceived, msg) {
		t.Fatalf("got %q, want %q", clientReceived, msg)
	}
}

// ---------------------------------------------------------------------------
// Encrypted round-trip
// ---------------------------------------------------------------------------

func TestIntegrationEncryptedRoundTrip(t *testing.T) {
	port := 17101
	key := "test-secret-key"

	server := NewServer(Config{Port: port, EncryptionKey: key, ConnectTimeout: 5 * time.Second})

	done := make(chan []byte, 1)
	server.OnData(func(id string, data []byte) {
		done <- data
	})

	go server.Listen()
	defer server.Close()
	time.Sleep(50 * time.Millisecond)

	client := NewClient(Config{Port: port})
	if err := client.Connect(); err != nil {
		t.Fatalf("connect: %v", err)
	}
	defer client.Disconnect()
	go client.Listen()
	time.Sleep(50 * time.Millisecond)

	msg := []byte("encrypted payload")
	client.SendEncrypted(msg, key)

	select {
	case got := <-done:
		if !bytes.Equal(got, msg) {
			t.Fatalf("got %q, want %q", got, msg)
		}
	case <-time.After(3 * time.Second):
		t.Fatal("timeout waiting for encrypted message")
	}
}

// ---------------------------------------------------------------------------
// Stream round-trip
// ---------------------------------------------------------------------------

func TestIntegrationStreamRoundTrip(t *testing.T) {
	port := 17102

	server := NewServer(Config{Port: port, ConnectTimeout: 5 * time.Second})

	done := make(chan []byte, 1)
	server.OnStream(func(id string, data []byte) {
		done <- data
	})

	go server.Listen()
	defer server.Close()
	time.Sleep(50 * time.Millisecond)

	client := NewClient(Config{Port: port, ChunkSize: 1024})
	if err := client.Connect(); err != nil {
		t.Fatalf("connect: %v", err)
	}
	defer client.Disconnect()
	go client.Listen()
	time.Sleep(50 * time.Millisecond)

	// 10 KB payload, 1 KB chunks = 10 chunks
	msg := make([]byte, 10*1024)
	for i := range msg {
		msg[i] = byte(i % 256)
	}
	client.SendStream(msg)

	select {
	case got := <-done:
		if !bytes.Equal(got, msg) {
			t.Fatalf("stream mismatch (got %d bytes, want %d)", len(got), len(msg))
		}
	case <-time.After(5 * time.Second):
		t.Fatal("timeout waiting for stream reassembly")
	}
}

// ---------------------------------------------------------------------------
// Multiple clients + broadcast
// ---------------------------------------------------------------------------

func TestIntegrationBroadcast(t *testing.T) {
	port := 17103

	server := NewServer(Config{Port: port, ConnectTimeout: 5 * time.Second})
	server.OnData(func(id string, data []byte) {
		server.BroadcastExcept(id, data)
	})

	go server.Listen()
	defer server.Close()
	time.Sleep(50 * time.Millisecond)

	// connect 3 clients
	clients := make([]*Client, 3)
	received := make([]chan []byte, 3)
	for i := 0; i < 3; i++ {
		clients[i] = NewClient(Config{Port: port})
		received[i] = make(chan []byte, 10)
		idx := i
		clients[i].OnData(func(data []byte) {
			received[idx] <- data
		})
		if err := clients[i].Connect(); err != nil {
			t.Fatalf("connect client %d: %v", i, err)
		}
		go clients[i].Listen()
	}
	defer func() {
		for _, c := range clients {
			c.Disconnect()
		}
	}()
	time.Sleep(100 * time.Millisecond)

	// client 0 sends a message - clients 1 and 2 should receive it
	msg := []byte("broadcast test")
	clients[0].Send(msg)

	for i := 1; i < 3; i++ {
		select {
		case got := <-received[i]:
			if !bytes.Equal(got, msg) {
				t.Fatalf("client %d got %q, want %q", i, got, msg)
			}
		case <-time.After(3 * time.Second):
			t.Fatalf("timeout waiting for client %d to receive broadcast", i)
		}
	}

	// client 0 should NOT receive its own broadcast
	select {
	case got := <-received[0]:
		t.Fatalf("client 0 should not receive its own message, got %q", got)
	case <-time.After(200 * time.Millisecond):
		// good - no message
	}
}

// ---------------------------------------------------------------------------
// MaxClients enforcement
// ---------------------------------------------------------------------------

func TestIntegrationMaxClients(t *testing.T) {
	port := 17104

	server := NewServer(Config{Port: port, MaxClients: 2, ConnectTimeout: 5 * time.Second})
	go server.Listen()
	defer server.Close()
	time.Sleep(50 * time.Millisecond)

	// connect 2 clients successfully
	c1 := NewClient(Config{Port: port})
	c2 := NewClient(Config{Port: port})
	if err := c1.Connect(); err != nil {
		t.Fatalf("c1 connect: %v", err)
	}
	go c1.Listen()
	if err := c2.Connect(); err != nil {
		t.Fatalf("c2 connect: %v", err)
	}
	go c2.Listen()
	time.Sleep(100 * time.Millisecond)

	if server.ClientCount() != 2 {
		t.Fatalf("expected 2 clients, got %d", server.ClientCount())
	}

	// third client should be rejected
	c3 := NewClient(Config{Port: port})
	rejected := make(chan string, 1)
	c3.OnReject(func(reason string) {
		rejected <- reason
	})
	if err := c3.Connect(); err != nil {
		t.Fatalf("c3 connect: %v", err)
	}
	go c3.Listen()

	select {
	case reason := <-rejected:
		if reason != "server full" {
			t.Fatalf("expected 'server full', got %q", reason)
		}
	case <-time.After(3 * time.Second):
		t.Fatal("timeout waiting for rejection")
	}

	c1.Disconnect()
	c2.Disconnect()
	c3.Disconnect()
}

// ---------------------------------------------------------------------------
// Graceful drain
// ---------------------------------------------------------------------------

func TestIntegrationDrain(t *testing.T) {
	port := 17105

	server := NewServer(Config{Port: port, DrainTimeout: 3 * time.Second, ConnectTimeout: 5 * time.Second})

	var clientID string
	server.OnConnect(func(id string) {
		clientID = id
	})

	go server.Listen()
	time.Sleep(50 * time.Millisecond)

	client := NewClient(Config{Port: port})
	if err := client.Connect(); err != nil {
		t.Fatalf("connect: %v", err)
	}
	go client.Listen()
	time.Sleep(100 * time.Millisecond)

	// start drain in background
	drained := make(chan int, 1)
	go func() {
		remaining := server.Drain()
		drained <- remaining
	}()

	// new connections should fail
	time.Sleep(100 * time.Millisecond)
	c2 := NewClient(Config{Port: port})
	err := c2.Connect()
	if err == nil {
		// connected but should be rejected immediately
		c2.Disconnect()
	}

	// disconnect the existing client - drain should complete
	client.Disconnect()

	select {
	case remaining := <-drained:
		if remaining != 0 {
			t.Fatalf("expected 0 remaining, got %d", remaining)
		}
	case <-time.After(5 * time.Second):
		t.Fatal("timeout waiting for drain")
	}

	_ = clientID // used to verify connect happened
}

// ---------------------------------------------------------------------------
// P2P encrypted mesh
// ---------------------------------------------------------------------------

func TestIntegrationP2PMesh(t *testing.T) {
	portA := 17110
	portB := 17111

	peerA := NewPeer(PeerConfig{Port: portA})
	peerB := NewPeer(PeerConfig{Port: portB})

	received := make(chan string, 1)
	peerB.OnData(func(peerID string, data []byte) {
		received <- string(data)
	})

	go peerA.Listen()
	go peerB.Listen()
	defer peerA.Close()
	defer peerB.Close()
	time.Sleep(100 * time.Millisecond)

	// B connects to A
	remotePeerID, err := peerB.Connect("127.0.0.1", portA)
	if err != nil {
		t.Fatalf("connect: %v", err)
	}
	time.Sleep(100 * time.Millisecond)

	// A sends to B
	msg := "hello p2p mesh"
	if err := peerA.SendTo(remotePeerID, []byte(msg)); err != nil {
		// peerA might know B by B's ID from the handshake
		// try broadcast instead
		peerA.Broadcast([]byte(msg))
	}

	select {
	case got := <-received:
		if got != msg {
			t.Fatalf("got %q, want %q", got, msg)
		}
	case <-time.After(3 * time.Second):
		t.Fatal("timeout waiting for P2P message")
	}

	if peerA.PeerCount() != 1 || peerB.PeerCount() != 1 {
		t.Fatalf("expected 1 peer each, got A=%d B=%d", peerA.PeerCount(), peerB.PeerCount())
	}

	fmt.Printf("P2P mesh test passed: peerA=%s, peerB=%s\n", peerA.GetPeerID()[:8], peerB.GetPeerID()[:8])
}