package notify

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

func TestClientPeerAttachRenamesAcceptedPeer(t *testing.T) {
	secret := []byte("0123456789abcdef0123456789abcdef")
	server := newRunningPeerAttachServerForTest(t, func(server *ServerCommon) {
		server.SetSecretKey(secret)
	})
	client := NewClient().(*ClientCommon)
	client.SetSecretKey(secret)

	left, right := net.Pipe()
	defer right.Close()

	accepted := bootstrapPeerAttachLogicalForTest(t, server, right)
	tempID := accepted.ClientID
	if err := client.ConnectByConn(left); err != nil {
		t.Fatalf("ConnectByConn failed: %v", err)
	}

	if got := server.GetLogicalConn(client.peerIdentity); got != accepted {
		t.Fatalf("stable peer lookup mismatch: got %v want %v", got, accepted)
	}
	if got := server.GetLogicalConn(tempID); got != nil {
		t.Fatalf("temporary accepted peer should be removed after attach, got %+v", got)
	}
	if got, want := accepted.ClientID, client.peerIdentity; got != want {
		t.Fatalf("accepted client id mismatch: got %q want %q", got, want)
	}

	client.setByeFromServer(true)
	if err := client.Stop(); err != nil {
		t.Fatalf("Stop failed: %v", err)
	}
}

func TestClientPeerAttachReusesExistingPeerOnTransportReattach(t *testing.T) {
	secret := []byte("0123456789abcdef0123456789abcdef")
	server := newRunningPeerAttachServerForTest(t, func(server *ServerCommon) {
		server.SetSecretKey(secret)
	})
	server.SetLink("echo", func(message *Message) {
		_ = message.Reply([]byte("pong"))
	})

	client := NewClient().(*ClientCommon)
	client.SetSecretKey(secret)

	firstLeft, firstRight := net.Pipe()
	defer firstRight.Close()
	bootstrapPeerAttachLogicalForTest(t, server, firstRight)
	if err := client.ConnectByConn(firstLeft); err != nil {
		t.Fatalf("initial ConnectByConn failed: %v", err)
	}

	stablePeer := server.GetLogicalConn(client.peerIdentity)
	if stablePeer == nil {
		t.Fatal("stable peer should exist after initial attach")
	}

	client2 := NewClient().(*ClientCommon)
	client2.SetSecretKey(secret)
	client2.peerIdentity = client.peerIdentity

	secondLeft, secondRight := net.Pipe()
	defer secondRight.Close()
	temp := bootstrapPeerAttachLogicalForTest(t, server, secondRight)
	tempID := temp.ClientID

	if err := client2.ConnectByConn(secondLeft); err != nil {
		t.Fatalf("second ConnectByConn failed: %v", err)
	}

	if got := server.GetLogicalConn(client2.peerIdentity); got != stablePeer {
		t.Fatalf("stable peer should be reused on handoff: got %v want %v", got, stablePeer)
	}
	if got := server.GetLogicalConn(tempID); got != nil {
		t.Fatalf("temporary peer should be removed after handoff, got %+v", got)
	}
	if got := stablePeer.clientConnTransportSnapshot(); got != secondRight {
		t.Fatalf("stable peer transport mismatch after handoff: got %v want %v", got, secondRight)
	}

	reply, err := client2.SendWait("echo", []byte("ping"), time.Second)
	if err != nil {
		t.Fatalf("SendWait after handoff failed: %v", err)
	}
	if got, want := string(reply.Value), "pong"; got != want {
		t.Fatalf("reply mismatch: got %q want %q", got, want)
	}

	client2.setByeFromServer(true)
	if err := client2.Stop(); err != nil {
		t.Fatalf("second client Stop failed: %v", err)
	}

	client.setByeFromServer(true)
	if err := client.Stop(); err != nil {
		t.Fatalf("first client Stop failed: %v", err)
	}
}

func TestReplyPeerAttachUsesInboundConnWithoutWaitingSignalAck(t *testing.T) {
	secret := []byte("0123456789abcdef0123456789abcdef")
	server := newRunningPeerAttachServerForTest(t, func(server *ServerCommon) {
		server.SetSecretKey(secret)
		if err := UseSignalReliabilityServer(server, &SignalReliabilityOptions{Enabled: true}); err != nil {
			t.Fatalf("UseSignalReliabilityServer failed: %v", err)
		}
	})

	clientConn, serverConn := net.Pipe()
	defer clientConn.Close()
	defer serverConn.Close()

	logical := bootstrapPeerAttachLogicalForTest(t, server, serverConn)
	originalProfile := logical.transportProtectionProfileSnapshot()
	message := Message{
		NetType:       NET_SERVER,
		LogicalConn:   logical,
		TransportConn: logical.CurrentTransportConn(),
		TransferMsg: TransferMsg{
			ID:   42,
			Key:  systemPeerAttachKey,
			Type: MSG_SYS_WAIT,
		},
		Time:        time.Now(),
		inboundConn: serverConn,
	}
	message = hydrateServerMessagePeerFields(message)

	alternate, err := deriveModernPSKProtectionProfile([]byte("notify-peer-attach-reply-alternate"), testModernPSKOptions(), ProtectionManaged)
	if err != nil {
		t.Fatalf("deriveModernPSKProtectionProfile failed: %v", err)
	}
	logical.applyTransportProtectionProfile(alternate)

	done := make(chan error, 1)
	go func() {
		done <- server.replyPeerAttach(logical, message, peerAttachResponse{
			PeerID:   "peer-test",
			Accepted: true,
		})
	}()

	env := readServerEnvelopeFromConnWithProfile(t, server, originalProfile, clientConn, time.Second)
	if env.Kind != EnvelopeSignal {
		t.Fatalf("reply envelope kind = %v, want %v", env.Kind, EnvelopeSignal)
	}

	select {
	case err := <-done:
		if err != nil {
			t.Fatalf("replyPeerAttach failed: %v", err)
		}
	case <-time.After(200 * time.Millisecond):
		t.Fatal("replyPeerAttach should finish without waiting for transport ack")
	}

	transfer, err := unwrapTransferMsgEnvelope(env, server.sequenceDe)
	if err != nil {
		t.Fatalf("unwrapTransferMsgEnvelope failed: %v", err)
	}
	if transfer.Type != MSG_SYS_REPLY {
		t.Fatalf("reply type = %v, want %v", transfer.Type, MSG_SYS_REPLY)
	}
	if transfer.ID != 42 {
		t.Fatalf("reply id = %d, want %d", transfer.ID, 42)
	}
	if transfer.Key != systemPeerAttachKey {
		t.Fatalf("reply key = %q, want %q", transfer.Key, systemPeerAttachKey)
	}
}