package peer

import (
	"bytes"
	"net/netip"
	"testing"
)

// ----------------------------------------------------------------------------

type testUDPPacket struct {
	Addr netip.AddrPort
	Data []byte
}

type testUDPAddrPortWriter struct {
	written []testUDPPacket
}

func (w *testUDPAddrPortWriter) WriteToUDPAddrPort(b []byte, addr netip.AddrPort) (int, error) {
	w.written = append(w.written, testUDPPacket{
		Addr: addr,
		Data: bytes.Clone(b),
	})
	return len(b), nil
}

func (w *testUDPAddrPortWriter) Written() []testUDPPacket {
	out := w.written
	w.written = []testUDPPacket{}
	return out
}

// ----------------------------------------------------------------------------

type testPacket string

func (p testPacket) Marshal(b []byte) []byte {
	b = b[:len(p)]
	copy(b, []byte(p))
	return b
}

// ----------------------------------------------------------------------------

func testConnWriter_getTestRoutes() (local, remote, relayLocal, relayRemote *peerRoute) {
	localKeys := generateKeys()
	remoteKeys := generateKeys()

	local = newPeerRoute(2)
	local.Up = true
	local.Relay = false
	local.PubSignKey = remoteKeys.PubSignKey
	local.ControlCipher = newControlCipher(localKeys.PrivKey, remoteKeys.PubKey)
	local.DataCipher = newDataCipher()
	local.RemoteAddr = netip.AddrPortFrom(netip.AddrFrom4([4]byte{1, 1, 1, 2}), 100)

	remote = newPeerRoute(1)
	remote.Up = true
	remote.Relay = false
	remote.PubSignKey = localKeys.PubSignKey
	remote.ControlCipher = newControlCipher(remoteKeys.PrivKey, localKeys.PubKey)
	remote.DataCipher = local.DataCipher
	remote.RemoteAddr = netip.AddrPortFrom(netip.AddrFrom4([4]byte{1, 1, 1, 1}), 100)

	rLocalKeys := generateKeys()
	rRemoteKeys := generateKeys()

	relayLocal = newPeerRoute(3)
	relayLocal.Up = true
	relayLocal.Relay = true
	relayLocal.Direct = true
	relayLocal.PubSignKey = rRemoteKeys.PubSignKey
	relayLocal.ControlCipher = newControlCipher(rLocalKeys.PrivKey, rRemoteKeys.PubKey)
	relayLocal.DataCipher = newDataCipher()
	relayLocal.RemoteAddr = netip.AddrPortFrom(netip.AddrFrom4([4]byte{1, 1, 1, 3}), 100)

	relayRemote = newPeerRoute(1)
	relayRemote.Up = true
	relayRemote.Relay = false
	relayRemote.Direct = true
	relayRemote.PubSignKey = rLocalKeys.PubSignKey
	relayRemote.ControlCipher = newControlCipher(rRemoteKeys.PrivKey, rLocalKeys.PubKey)
	relayRemote.DataCipher = relayLocal.DataCipher
	relayRemote.RemoteAddr = netip.AddrPortFrom(netip.AddrFrom4([4]byte{1, 1, 1, 1}), 100)

	return
}

// ----------------------------------------------------------------------------

// Testing if we can send a control packet directly to the remote route.
func TestConnWriter_SendControlPacket_direct(t *testing.T) {
	route, rRoute, _, _ := testConnWriter_getTestRoutes()
	route.Direct = true

	writer := &testUDPAddrPortWriter{}
	w := newConnWriter(writer, rRoute.IP)
	in := testPacket("hello world!")

	w.SendControlPacket(in, route)
	out := writer.Written()
	if len(out) != 1 {
		t.Fatal(out)
	}

	if out[0].Addr != route.RemoteAddr {
		t.Fatal(out[0])
	}

	dec, ok := rRoute.ControlCipher.Decrypt(out[0].Data, make([]byte, 1024))
	if !ok {
		t.Fatal(ok)
	}
	if string(dec) != string(in) {
		t.Fatal(dec)
	}
}

// Testing if we can relay a packet via an intermediary.
func TestConnWriter_RelayControlPacket_relay(t *testing.T) {
	route, rRoute, relay, rRelay := testConnWriter_getTestRoutes()

	writer := &testUDPAddrPortWriter{}
	w := newConnWriter(writer, rRoute.IP)
	in := testPacket("hello world!")

	w.RelayControlPacket(in, route, relay)

	out := writer.Written()
	if len(out) != 1 {
		t.Fatal(out)
	}

	if out[0].Addr != relay.RemoteAddr {
		t.Fatal(out[0])
	}

	dec, ok := rRelay.DataCipher.Decrypt(out[0].Data, make([]byte, 1024))
	if !ok {
		t.Fatal(ok)
	}

	dec2, ok := rRoute.ControlCipher.Decrypt(dec, make([]byte, 1024))
	if !ok {
		t.Fatal(ok)
	}

	if string(dec2) != string(in) {
		t.Fatal(dec2)
	}
}

// Testing that we can send a data packet directly to a remote route.
func TestConnWriter_SendDataPacket_direct(t *testing.T) {
	route, rRoute, _, _ := testConnWriter_getTestRoutes()
	route.Direct = true

	writer := &testUDPAddrPortWriter{}
	w := newConnWriter(writer, rRoute.IP)

	in := []byte("hello world!")
	w.SendDataPacket(in, route)

	out := writer.Written()
	if len(out) != 1 {
		t.Fatal(out)
	}

	if out[0].Addr != route.RemoteAddr {
		t.Fatal(out[0])
	}

	dec, ok := rRoute.DataCipher.Decrypt(out[0].Data, make([]byte, 1024))
	if !ok {
		t.Fatal(ok)
	}

	if !bytes.Equal(dec, in) {
		t.Fatal(dec)
	}
}

// Testing that we can relay a data packet via a relay.
func TestConnWriter_RelayDataPacket_relay(t *testing.T) {
	route, rRoute, relay, rRelay := testConnWriter_getTestRoutes()

	writer := &testUDPAddrPortWriter{}
	w := newConnWriter(writer, rRoute.IP)
	in := []byte("Hello world!")

	w.RelayDataPacket(in, route, relay)

	out := writer.Written()
	if len(out) != 1 {
		t.Fatal(out)
	}

	if out[0].Addr != relay.RemoteAddr {
		t.Fatal(out[0])
	}

	dec, ok := rRelay.DataCipher.Decrypt(out[0].Data, make([]byte, 1024))
	if !ok {
		t.Fatal(ok)
	}

	dec2, ok := rRoute.DataCipher.Decrypt(dec, make([]byte, 1024))
	if !ok {
		t.Fatal(ok)
	}

	if !bytes.Equal(dec2, in) {
		t.Fatal(dec2)
	}
}

// Testing that we can send an already encrypted packet.
func TestConnWriter_SendEncryptedDataPacket(t *testing.T) {
	route, rRoute, _, _ := testConnWriter_getTestRoutes()

	writer := &testUDPAddrPortWriter{}
	w := newConnWriter(writer, rRoute.IP)
	in := []byte("Hello world!")

	w.SendEncryptedDataPacket(in, route)

	out := writer.Written()
	if len(out) != 1 {
		t.Fatal(out)
	}

	if out[0].Addr != route.RemoteAddr {
		t.Fatal(out[0])
	}

	if !bytes.Equal(out[0].Data, in) {
		t.Fatal(out[0])
	}
}