package kcp

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

// TestKcpLoopback runs two KCP instances over real UDP sockets and exchanges
// fragmented messages in both directions.
func TestKcpLoopback(t *testing.T) {
	connA, err := net.ListenUDP("udp", &net.UDPAddr{IP: net.IPv4(127, 0, 0, 1)})
	if err != nil {
		t.Fatal(err)
	}
	connB, err := net.ListenUDP("udp", &net.UDPAddr{IP: net.IPv4(127, 0, 0, 1)})
	if err != nil {
		t.Fatal(err)
	}

	addrB := connB.LocalAddr().(*net.UDPAddr)
	addrA := connA.LocalAddr().(*net.UDPAddr)

	var now uint32 = 1000

	kcpA := New(77, 1234, false, func(data []byte) error {
		_, err := connA.WriteToUDP(data, addrB)
		return err
	})
	kcpB := New(77, 1234, false, func(data []byte) error {
		_, err := connB.WriteToUDP(data, addrA)
		return err
	})

	pump := func() {
		bufA := make([]byte, 65536)
		bufB := make([]byte, 65536)
		_ = connB.SetReadDeadline(time.Now().Add(50 * time.Millisecond))
		_ = connA.SetReadDeadline(time.Now().Add(50 * time.Millisecond))
		for {
			_ = connB.SetReadDeadline(time.Now().Add(10 * time.Millisecond))
			n, _, err := connB.ReadFrom(bufB)
			if err != nil {
				break
			}
			if err := kcpB.Input(bufB[:n]); err != nil {
				t.Fatalf("B input: %v", err)
			}
		}
		for {
			_ = connA.SetReadDeadline(time.Now().Add(10 * time.Millisecond))
			n, _, err := connA.ReadFrom(bufA)
			if err != nil {
				break
			}
			if err := kcpA.Input(bufA[:n]); err != nil {
				t.Fatalf("A input: %v", err)
			}
		}
	}

	drain := func(k *KCP) [][]byte {
		var out [][]byte
		buf := make([]byte, 65536)
		for {
			n, err := k.Recv(buf)
			if err != nil {
				break
			}
			msg := make([]byte, n)
			copy(msg, buf[:n])
			out = append(out, msg)
		}
		return out
	}

	tick := func() {
		now += Interval
		if err := kcpA.Update(now); err != nil {
			t.Fatal(err)
		}
		if err := kcpB.Update(now); err != nil {
			t.Fatal(err)
		}
	}

	if err := kcpA.Update(now); err != nil {
		t.Fatal(err)
	}
	if err := kcpB.Update(now); err != nil {
		t.Fatal(err)
	}

	// large message -> fragmented; small one -> single segment
	msg1 := make([]byte, 20000)
	for i := range msg1 {
		msg1[i] = byte(i % 251)
	}
	msg2 := []byte("hello hsr")

	if _, err := kcpA.Send(msg1); err != nil {
		t.Fatal(err)
	}
	if _, err := kcpA.Send(msg2); err != nil {
		t.Fatal(err)
	}

	delivered := 0
	for i := 0; i < 500 && delivered < 2; i++ {
		tick()
		pump()
		delivered += len(drain(kcpB))
	}
	if delivered != 2 {
		t.Fatalf("expected 2 messages delivered, got %d", delivered)
	}

	// verify payload integrity with a retransmission-tolerant final exchange
	if _, err := kcpA.Send([]byte("final")); err != nil {
		t.Fatal(err)
	}
	deadline := time.Now().Add(2 * time.Second)
	for time.Now().Before(deadline) {
		tick()
		pump()
		msgs := drain(kcpB)
		if len(msgs) > 0 {
			if string(msgs[0]) != "final" {
				t.Fatalf("payload mismatch: %q", msgs[0])
			}
			return
		}
	}
	t.Fatal("timed out waiting for final message")
}