package server

import (
	"net"
	"strings"
	"sync"
	"testing"
	"time"

	kcp "github.com/xtaci/kcp-go/v5"

	"omnisocketgo/cmd/internal/latencylog"
	"omnisocketgo/cmd/internal/protocol"
	"omnisocketgo/cmd/internal/transport"
)

// TestUDPRelayKCPForwardAndReturn 验证 KCP 通过 UDP relay 的完整双向转发路径：
// peer-b -> D(KCP hub) -> C(UDP relay) -> peer-a 以及反向。
func TestUDPRelayKCPForwardAndReturn(t *testing.T) {
	// 启动 D（KCP Hub）
	hub, hubAddr, hubCleanup := startKCPHubForRelay(t)
	defer hubCleanup()

	// 启动 C（UDP Relay），upstream 指向 D
	relayAddr := startUDPRelay(t, hubAddr)

	// peer-b 直连 D（KCP）
	peerBConn := dialKCPPeer(t, hubAddr)
	// peer-a 连 C（通过 relay 间接连到 D）
	peerAConn := dialKCPPeer(t, relayAddr)

	// 注册 peer-b
	if err := peerBConn.Send(protocol.Message{
		Type: protocol.MessageTypeRegister,
		From: "peer-b",
		To:   protocol.ServerPeerID,
	}); err != nil {
		t.Fatalf("peerB register: %v", err)
	}

	// 注册 peer-a（通过 relay）
	if err := peerAConn.Send(protocol.Message{
		Type: protocol.MessageTypeRegister,
		From: "peer-a",
		To:   protocol.ServerPeerID,
	}); err != nil {
		t.Fatalf("peerA register: %v", err)
	}

	waitForRelay(t, func() bool {
		return hub.HasPeer("peer-a") && hub.HasPeer("peer-b")
	}, "both peers to be registered")

	// peer-b -> peer-a（路径: B -> D -> C -> A）
	if err := peerBConn.Send(protocol.Message{
		Type: protocol.MessageTypeText,
		ID:   1,
		From: "peer-b",
		To:   "peer-a",
		Body: []byte("hello from peer-b"),
	}); err != nil {
		t.Fatalf("peerB send text: %v", err)
	}

	msg, err := peerAConn.Receive()
	if err != nil {
		t.Fatalf("peerA receive: %v", err)
	}
	if msg.Type != protocol.MessageTypeText {
		t.Fatalf("message type = %s, want text", msg.Type)
	}
	if msg.From != "peer-b" {
		t.Fatalf("message from = %s, want peer-b", msg.From)
	}
	if string(msg.Body) != "hello from peer-b" {
		t.Fatalf("message body = %q, want %q", string(msg.Body), "hello from peer-b")
	}

	// peer-a -> peer-b（路径: A -> C -> D -> B）
	if err := peerAConn.Send(protocol.Message{
		Type: protocol.MessageTypeText,
		ID:   2,
		From: "peer-a",
		To:   "peer-b",
		Body: []byte("reply from peer-单个 downstream peer 通过 relay 连到 KCP server”这条
  链路是成立的，转发逻辑本身没有明显的地址错误。cmd/internal/server/udp_relay.go 里就是原
  样双向转发，下游来的包会记录 clientAddr 并写给上游，上游回来的包再写回这个 clientAddr。
  关键代码在 cmd/internal/server/udp_relay.go:68 和 cmd/internal/server/udp_relay.go:89。

  还有一个关键事实：kcppeer 里那句 connected to ... as ... (KCP) 不能证明 peer-a 真的在
  hub 注册成功a"),
	}); err != nil {
		t.Fatalf("peerA send text: %v", err)
	}

	msg2, err := peerBConn.Receive()
	if err != nil {
		t.Fatalf("peerB receive: %v", err)
	}
	if msg2.Type != protocol.MessageTypeText {
		t.Fatalf("message type = %s, want text", msg2.Type)
	}
	if msg2.From != "peer-a" {
		t.Fatalf("message from = %s, want peer-a", msg2.From)
	}
	if string(msg2.Body) != "reply from peer-a" {
		t.Fatalf("message body = %q, want %q", string(msg2.Body), "reply from peer-a")
	}
}

// TestUDPRelayKCPFileMessage 验证通过 relay 转发 KCP 文件消息。
func TestUDPRelayKCPFileMessage(t *testing.T) {
	hub, hubAddr, hubCleanup := startKCPHubForRelay(t)
	defer hubCleanup()

	relayAddr := startUDPRelay(t, hubAddr)

	peerBConn := dialKCPPeer(t, hubAddr)
	peerAConn := dialKCPPeer(t, relayAddr)

	_ = peerBConn.Send(protocol.Message{
		Type: protocol.MessageTypeRegister,
		From: "peer-b",
		To:   protocol.ServerPeerID,
	})
	_ = peerAConn.Send(protocol.Message{
		Type: protocol.MessageTypeRegister,
		From: "peer-a",
		To:   protocol.ServerPeerID,
	})

	waitForRelay(t, func() bool {
		return hub.HasPeer("peer-a") && hub.HasPeer("peer-b")
	}, "both peers to be registered")

	if err := peerBConn.Send(protocol.Message{
		Type:     protocol.MessageTypeFile,
		ID:       1,
		From:     "peer-b",
		To:       "peer-a",
		FileName: "test.bin",
		Body:     []byte{0xDE, 0xAD, 0xBE, 0xEF},
	}); err != nil {
		t.Fatalf("peerB send file: %v", err)
	}

	msg, err := peerAConn.Receive()
	if err != nil {
		t.Fatalf("peerA receive: %v", err)
	}
	if msg.Type != protocol.MessageTypeFile {
		t.Fatalf("message type = %s, want file", msg.Type)
	}
	if msg.FileName != "test.bin" {
		t.Fatalf("file name = %q, want %q", msg.FileName, "test.bin")
	}
	if len(msg.Body) != 4 || msg.Body[0] != 0xDE {
		t.Fatalf("file body mismatch: 单个 downstream peer 通过 relay 连到 KCP server”这条
  链路是成立的，转发逻辑本身没有明显的地址错误。cmd/internal/server/udp_relay.go 里就是原
  样双向转发，下游来的包会记录 clientAddr 并写给上游，上游回来的包再写回这个 clientAddr。
  关键代码在 cmd/internal/server/udp_relay.go:68 和 cmd/internal/server/udp_relay.go:89。

  还有一个关键事实：kcppeer 里那句 connected to ... as ... (KCP) 不能证明 peer-a 真的在
  hub 注册成功got %v", msg.Body)
	}
}

// startKCPHubForRelay 启动一个 KCP hub server，返回 hub、监听地址和 cleanup 函数。
func startKCPHubForRelay(t *testing.T) (*KCPHub, string, func()) {
	t.Helper()

	hub := NewKCPHub()

	listener, packetConn, err := transport.ListenKCPSessions("127.0.0.1:0", "", nil, latencylog.NodeRoleServer, "hub")
	if err != nil {
		t.Fatalf("ListenKCPSessions() error = %v", err)
	}

	var (
		wg   sync.WaitGroup
		stop = make(chan struct{})
	)

	wg.Add(1)
	go func() {
		defer wg.Done()
		for {
			session, acceptErr := listener.AcceptKCP()
			if acceptErr != nil {
				select {
				case <-stop:
					return
				default:
				}
				if strings.Contains(acceptErr.Error(), "closed") {
					return
				}
				t.Errorf("AcceptKCP() error = %v", acceptErr)
				return
			}

			wg.Add(1)
			go func(sess *kcp.UDPSession) {
				defer wg.Done()
				if serveErr := hub.ServeSession(sess); serveErr != nil {
					msg := serveErr.Error()
					if !strings.Contains(msg, "closed") && !strings.Contains(msg, "broken pipe") {
						t.Logf("hub.ServeSession() ended with %v", serveErr)
					}
				}
			}(session)
		}
	}()

	cleanup := func() {
		close(stop)
		_ = listener.Close()
		_ = packetConn.Close()
		wg.Wait()
	}

	return hub, listener.Addr().String(), cleanup
}

// dialKCPPeer 创建一条到指定地址的 KCP 连接，用于测试。
func dialKCPPeer(t *testing.T, serverAddr string) *transport.KCPConn {
	t.Helper()

	session, err := transport.DialKCPSession(serverAddr, "", "", nil, latencylog.NodeRolePeer, "test")
	if err != nil {
		t.Fatalf("DialKCPSession(%s) error = %v", serverAddr, err)
	}

	conn, err := transport.NewKCPConn(session)
	if err != nil {
		_ = session.Close()
		t.Fatalf("NewKCPConn() error = %v", err)
	}

	t.Cleanup(func() {
		_ = conn.Close()
	})

	return conn
}

// startUDPRelay 创建并启动一个 UDPRelay，返回其监听地址字符串。
func startUDPRelay(t *testing.T, upstreamAddr string) string {
	t.Helper()

	addr, err := net.ResolveUDPAddr("udp", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("ResolveUDPAddr() error = %v", err)
	}

	conn, err := net.ListenUDP("udp", addr)
	if err != nil {
		t.Fatalf("ListenUDP() error = %v", err)
	}

	relay, err := NewUDPRelay(conn, upstreamAddr)
	if err != nil {
		_ = conn.Close()
		t.Fatalf("NewUDPRelay() error = %v", err)
	}

	go func() {
		_ = relay.Serve()
	}()

	t.Cleanup(func() {
		_ = relay.Close()
	})

	return conn.LocalAddr().String()
}

// waitForRelay 轮询等待条件满足，超时则 fail。
func waitForRelay(t *testing.T, condition func() bool, description string) {
	t.Helper()

	deadline := time.Now().Add(2 * time.Second)
	for time.Now().Before(deadline) {
		if condition() {
			return
		}
		time.Sleep(10 * time.Millisecond)
	}

	t.Fatalf("timed out waiting for %s", description)
}