package notify

import (
	"b612.me/stario"
	"net"
	"testing"
	"time"
)

func TestClientPeerAttachRenamesAcceptedPeer(t *testing.T) {
	secret := []byte("0123456789abcdef0123456789abcdef")
	server := newRunningPeerAttachServerForTest(t, func(server *ServerCommon) {
		server.SetSecretKey(secret)
		if err := UseSignalReliabilityServer(server, &SignalReliabilityOptions{
			Enabled:    true,
			AckTimeout: 200 * time.Millisecond,
			SendRetry:  2,
		}); err != nil {
			t.Fatal(err)
		}
	})
	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)
	}
}

func TestReplyPeerAttachUsesCapturedProfileWithoutInboundConn(t *testing.T) {
	secret := []byte("0123456789abcdef0123456789abcdef")
	server := newRunningPeerAttachServerForTest(t, func(server *ServerCommon) {
		server.SetSecretKey(secret)
		if err := UseSignalReliabilityServer(server, &SignalReliabilityOptions{
			Enabled:    true,
			AckTimeout: 2 * time.Second,
			SendRetry:  2,
		}); err != nil {
			t.Fatal(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:   43,
			Key:  systemPeerAttachKey,
			Type: MSG_SYS_WAIT,
		},
		Time:                    time.Now(),
		inboundTransportProfile: &originalProfile,
	}

	alternate, err := deriveModernPSKProtectionProfile([]byte("notify-peer-attach-no-inbound-conn"), testModernPSKOptions(), ProtectionManaged)
	if err != nil {
		t.Fatalf("deriveModernPSKProtectionProfile failed: %v", err)
	}
	logical.applyTransportProtectionProfile(alternate)
	transition := logical.installInboundTransitionProfile(originalProfile)
	defer logical.clearInboundTransitionProfile(transition)

	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)
	}
	ackPlain, err := server.encodeEnvelopePlain(newSignalAckEnvelope(env.ID))
	if err != nil {
		t.Fatal(err)
	}
	ackPayload, err := encryptTransportPayloadCodec(originalProfile.mode, originalProfile.runtime, originalProfile.msgEn, originalProfile.secretKey, ackPlain)
	if err != nil {
		t.Fatal(err)
	}
	if err := writeFullToConn(clientConn, stario.NewQueue().BuildMessage(ackPayload)); err != nil {
		t.Fatalf("write bootstrap signal ack: %v", err)
	}
	select {
	case err := <-done:
		if err != nil {
			t.Fatalf("replyPeerAttach failed: %v", err)
		}
	case <-time.After(time.Second):
		t.Fatal("replyPeerAttach should finish without an inbound stream conn")
	}
}

func TestPeerAttachTransitionProfileDecryptsBootstrapFramesAndClears(t *testing.T) {
	bootstrap, err := deriveModernPSKProtectionProfile([]byte("notify-peer-transition-bootstrap"), testModernPSKOptions(), ProtectionManaged)
	if err != nil {
		t.Fatal(err)
	}
	steady, err := deriveModernPSKProtectionProfile([]byte("notify-peer-transition-steady"), testModernPSKOptions(), ProtectionManaged)
	if err != nil {
		t.Fatal(err)
	}
	payload, err := encryptTransportPayloadCodec(bootstrap.mode, bootstrap.runtime, bootstrap.msgEn, bootstrap.secretKey, []byte("bootstrap-frame"))
	if err != nil {
		t.Fatal(err)
	}

	t.Run("client", func(t *testing.T) {
		client := NewClient().(*ClientCommon)
		client.setClientTransportProtectionProfile(steady)
		transition := client.installInboundTransitionProfile(bootstrap)
		plain, release, err := client.decryptTransportPayloadPooled(append([]byte(nil), payload...), nil)
		if release != nil {
			release()
		}
		if err != nil || string(plain) != "bootstrap-frame" {
			t.Fatalf("client transition decrypt plain=%q err=%v", plain, err)
		}
		client.clearInboundTransitionProfile(transition)
		if _, _, err := client.decryptTransportPayloadPooled(append([]byte(nil), payload...), nil); err == nil {
			t.Fatal("client accepted bootstrap frame after transition fallback cleared")
		}
	})

	t.Run("server", func(t *testing.T) {
		logical := newServerLogicalConn(nil, "transition-server", nil)
		logical.applyTransportProtectionProfile(steady)
		transition := logical.installInboundTransitionProfile(bootstrap)
		plain, release, err := (&ServerCommon{}).decryptTransportPayloadLogicalPooled(logical, append([]byte(nil), payload...), nil)
		if release != nil {
			release()
		}
		if err != nil || string(plain) != "bootstrap-frame" {
			t.Fatalf("server transition decrypt plain=%q err=%v", plain, err)
		}
		logical.clearInboundTransitionProfile(transition)
		if _, _, err := (&ServerCommon{}).decryptTransportPayloadLogicalPooled(logical, append([]byte(nil), payload...), nil); err == nil {
			t.Fatal("server accepted bootstrap frame after transition fallback cleared")
		}
	})
}