package node

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 = &peerRoute{
		IP:            2,
		Up:            true,
		Relay:         false,
		PubSignKey:    remoteKeys.PubSignKey,
		ControlCipher: newControlCipher(localKeys.PrivKey, remoteKeys.PubKey),
		DataCipher:    newDataCipher(),
		RemoteAddr:    netip.AddrPortFrom(netip.AddrFrom4([4]byte{1, 1, 1, 2}), 100),
	}

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

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

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

	relayRemote = &peerRoute{
		IP:            1,
		Up:            true,
		Relay:         false,
		Direct:        true,
		PubSignKey:    rLocalKeys.PubSignKey,
		ControlCipher: newControlCipher(rRemoteKeys.PrivKey, rLocalKeys.PubKey),
		DataCipher:    relayLocal.DataCipher,
		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_SendControlPacket_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 a nil relay doesn't cause an issue.
func TestConnWriter_SendControlPacket_relay_relayNil(t *testing.T) {
	route, rRoute, _, _ := testConnWriter_getTestRoutes()

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

	w.RelayControlPacket(in, route, nil)

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

}

// Testing that we don't send anything if the relay isn't up.
func TestConnWriter_SendControlPacket_relay_relayNotUp(t *testing.T) {
	route, rRoute, relay, _ := testConnWriter_getTestRoutes()
	relay.Up = false

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

	w.RelayControlPacket(in, route, relay)

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

// 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, nil)

	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_SendDataPacket_relay(t *testing.T) {
	route, rRoute, relay, rRelay := testConnWriter_getTestRoutes()

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

	w.SendDataPacket(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 don't attempt to relay if the relay is nil.
func TestConnWriter_SendDataPacket_relay_relayNil(t *testing.T) {
	route, rRoute, _, _ := testConnWriter_getTestRoutes()

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

	w.SendDataPacket(in, route, nil)

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

// Testing that we don't attempt to relay if the relay isn't up.
func TestConnWriter_SendDataPacket_relay_relayNotUp(t *testing.T) {
	route, rRoute, relay, _ := testConnWriter_getTestRoutes()
	relay.Up = false

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

	w.SendDataPacket(in, route, relay)

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