package peer

import (
	"bytes"
	"crypto/rand"
	"net/netip"
	"reflect"
	"sync/atomic"
	"testing"
)

type mockIfWriter struct {
	Written [][]byte
}

func (w *mockIfWriter) Write(b []byte) (int, error) {
	w.Written = append(w.Written, bytes.Clone(b))
	return len(b), nil
}

type mockEncryptedPacket struct {
	Packet []byte
	Route  *RemotePeer
}

type mockEncryptedPacketSender struct {
	Sent []mockEncryptedPacket
}

func (m *mockEncryptedPacketSender) SendEncryptedDataPacket(pkt []byte, route *RemotePeer) {
	m.Sent = append(m.Sent, mockEncryptedPacket{
		Packet: bytes.Clone(pkt),
		Route:  route,
	})
}

type mockControlMsgHandler struct {
	Messages []any
}

func (m *mockControlMsgHandler) HandleControlMsg(pkt any) {
	m.Messages = append(m.Messages, pkt)
}

type udpPipe struct {
	packets chan []byte
}

func newUDPPipe() *udpPipe {
	return &udpPipe{make(chan []byte, 1024)}
}

func (p *udpPipe) WriteToUDPAddrPort(b []byte, addr netip.AddrPort) (int, error) {
	p.packets <- bytes.Clone(b)
	return len(b), nil
}

func (p *udpPipe) ReadFromUDPAddrPort(b []byte) (n int, addr netip.AddrPort, err error) {
	packet := <-p.packets
	copy(b, packet)
	return len(packet), netip.AddrPort{}, nil
}

type connReaderTestHarness struct {
	Pipe         *udpPipe
	R            *connReader
	WRemote      *connWriter
	WRelayRemote *connWriter
	Remote       *RemotePeer
	RelayRemote  *RemotePeer
	IFace        *mockIfWriter
	Sender       *mockEncryptedPacketSender
	Super        *mockControlMsgHandler
}

// Peer 2 is indirect, peer 3 is direct.
func newConnReadeTestHarness() (h connReaderTestHarness) {
	pipe := newUDPPipe()
	routes := [256]*atomic.Pointer[RemotePeer]{}
	for i := range routes {
		routes[i] = &atomic.Pointer[RemotePeer]{}
		routes[i].Store(&RemotePeer{})
	}

	local, remote, relayLocal, relayRemote := testConnWriter_getTestRoutes()
	routes[2].Store(local)
	routes[3].Store(relayLocal)

	h.Pipe = pipe
	h.WRemote = newConnWriter(pipe, 2)
	h.WRelayRemote = newConnWriter(pipe, 3)

	h.Remote = remote
	h.RelayRemote = relayRemote
	h.IFace = &mockIfWriter{}
	h.Sender = &mockEncryptedPacketSender{}
	h.Super = &mockControlMsgHandler{}
	h.R = newConnReader(
		pipe,
		h.IFace,
		h.Sender,
		h.Super,
		1,
		routes)
	return h
}

// Testing that we can receive a control packet.
func TestConnReader_handleControlPacket(t *testing.T) {
	h := newConnReadeTestHarness()

	pkt := synPacket{TraceID: 1234}

	h.WRemote.SendControlPacket(pkt, h.Remote)

	h.R.handleNextPacket()

	if len(h.Super.Messages) != 1 {
		t.Fatal(h.Super.Messages)
	}

	msg := h.Super.Messages[0].(controlMsg[synPacket])
	if !reflect.DeepEqual(pkt, msg.Packet) {
		t.Fatal(msg.Packet)
	}
}

// Testing that a short packet is ignored.
func TestConnReader_handleNextPacket_short(t *testing.T) {
	h := newConnReadeTestHarness()

	h.Pipe.WriteToUDPAddrPort([]byte{1, 2, 3}, netip.AddrPort{})
	h.R.handleNextPacket()

	if len(h.Super.Messages) != 0 {
		t.Fatal(h.Super.Messages)
	}
}

// Testing that a packet with an unexpected stream ID is ignored.
func TestConnReader_handleNextPacket_unknownStreamID(t *testing.T) {
	h := newConnReadeTestHarness()

	pkt := synPacket{TraceID: 1234}

	encrypted := encryptControlPacket(1, h.Remote, pkt, newBuf(), newBuf())
	var header header
	header.Parse(encrypted)
	header.StreamID = 100
	header.Marshal(encrypted)

	h.WRemote.writeTo(encrypted, netip.AddrPort{})
	h.R.handleNextPacket()
	if len(h.Super.Messages) != 0 {
		t.Fatal(h.Super.Messages)
	}
}

// Testing that control packet without matching control cipher is ignored.
func TestConnReader_handleControlPacket_noCipher(t *testing.T) {
	h := newConnReadeTestHarness()

	pkt := synPacket{TraceID: 1234}

	//encrypted := h.WRemote.encryptControlPacket(pkt, h.Remote)
	encrypted := encryptControlPacket(1, h.Remote, pkt, newBuf(), newBuf())
	var header header
	header.Parse(encrypted)
	header.SourceIP = 10
	header.Marshal(encrypted)

	h.WRemote.writeTo(encrypted, netip.AddrPort{})
	h.R.handleNextPacket()
	if len(h.Super.Messages) != 0 {
		t.Fatal(h.Super.Messages)
	}
}

// Testing that control packet with incrrect destination IP is ignored.
func TestConnReader_handleControlPacket_incorrectDest(t *testing.T) {
	h := newConnReadeTestHarness()

	pkt := synPacket{TraceID: 1234}

	encrypted := encryptControlPacket(2, h.Remote, pkt, newBuf(), newBuf())
	var header header
	header.Parse(encrypted)
	header.DestIP++
	header.Marshal(encrypted)

	h.WRemote.writeTo(encrypted, netip.AddrPort{})
	h.R.handleNextPacket()
	if len(h.Super.Messages) != 0 {
		t.Fatal(h.Super.Messages)
	}
}

// Testing that modified control packet is ignored.
func TestConnReader_handleControlPacket_modified(t *testing.T) {
	h := newConnReadeTestHarness()

	pkt := synPacket{TraceID: 1234}

	encrypted := encryptControlPacket(2, h.Remote, pkt, newBuf(), newBuf())
	encrypted[len(encrypted)-1]++

	h.WRemote.writeTo(encrypted, netip.AddrPort{})
	h.R.handleNextPacket()
	if len(h.Super.Messages) != 0 {
		t.Fatal(h.Super.Messages)
	}
}

type unknownPacket struct{}

func (p unknownPacket) Marshal(buf []byte) []byte {
	buf = buf[:1]
	buf[0] = 100
	return buf
}

// Testing that an empty control packet is ignored.
func TestConnReader_handleControlPacket_unknownPacketType(t *testing.T) {
	h := newConnReadeTestHarness()

	pkt := unknownPacket{}

	encrypted := encryptControlPacket(2, h.Remote, pkt, newBuf(), newBuf())
	h.WRemote.writeTo(encrypted, netip.AddrPort{})
	h.R.handleNextPacket()
	if len(h.Super.Messages) != 0 {
		t.Fatal(h.Super.Messages)
	}
}

// Testing that a duplicate control packet is ignored.
func TestConnReader_handleControlPacket_duplicate(t *testing.T) {
	h := newConnReadeTestHarness()

	pkt := ackPacket{TraceID: 1234}

	h.WRemote.SendControlPacket(pkt, h.Remote)
	*h.Remote.Counter = *h.Remote.Counter - 1
	h.WRemote.SendControlPacket(pkt, h.Remote)

	h.R.handleNextPacket()
	h.R.handleNextPacket()

	if len(h.Super.Messages) != 1 {
		t.Fatal(h.Super.Messages)
	}

	msg := h.Super.Messages[0].(controlMsg[ackPacket])
	if !reflect.DeepEqual(pkt, msg.Packet) {
		t.Fatal(msg.Packet)
	}
}

// Testing that we can receive a data packet.
func TestConnReader_handleDataPacket(t *testing.T) {
	h := newConnReadeTestHarness()

	pkt := make([]byte, 1024)
	rand.Read(pkt)

	h.WRemote.SendDataPacket(pkt, h.Remote)

	h.R.handleNextPacket()

	if len(h.IFace.Written) != 1 {
		t.Fatal(h.IFace.Written)
	}

	if !bytes.Equal(pkt, h.IFace.Written[0]) {
		t.Fatal(h.IFace.Written)
	}
}

// Testing that data packet is ignored if route isn't up.
func TestConnReader_handleDataPacket_routeDown(t *testing.T) {
	h := newConnReadeTestHarness()

	pkt := make([]byte, 1024)
	rand.Read(pkt)

	h.WRemote.SendDataPacket(pkt, h.Remote)
	route := h.R.peers[2].Load()
	route.Up = false

	h.R.handleNextPacket()

	if len(h.IFace.Written) != 0 {
		t.Fatal(h.IFace.Written)
	}
}

// Testing that a duplicate data packet is ignored.
func TestConnReader_handleDataPacket_duplicate(t *testing.T) {
	h := newConnReadeTestHarness()

	pkt := make([]byte, 123)

	h.WRemote.SendDataPacket(pkt, h.Remote)
	*h.Remote.Counter = *h.Remote.Counter - 1
	h.WRemote.SendDataPacket(pkt, h.Remote)

	h.R.handleNextPacket()
	h.R.handleNextPacket()

	if len(h.IFace.Written) != 1 {
		t.Fatal(h.IFace.Written)
	}

	if !bytes.Equal(pkt, h.IFace.Written[0]) {
		t.Fatal(h.IFace.Written)
	}
}

// Testing that we can relay a data packet.
func TestConnReader_handleDataPacket_relay(t *testing.T) {
	h := newConnReadeTestHarness()

	pkt := make([]byte, 1024)
	rand.Read(pkt)

	h.RelayRemote.IP = 3
	h.WRemote.RelayDataPacket(pkt, h.RelayRemote, h.Remote)
	h.R.handleNextPacket()

	if len(h.Sender.Sent) != 1 {
		t.Fatal(h.Sender.Sent)
	}

}

// Testing that we drop a relayed packet if destination is down.
func TestConnReader_handleDataPacket_relayDown(t *testing.T) {
	h := newConnReadeTestHarness()

	pkt := make([]byte, 1024)
	rand.Read(pkt)

	h.RelayRemote.IP = 3
	relay := h.R.peers[3].Load()
	relay.Up = false

	h.WRemote.RelayDataPacket(pkt, h.RelayRemote, h.Remote)
	h.R.handleNextPacket()

	if len(h.Sender.Sent) != 0 {
		t.Fatal(h.Sender.Sent)
	}
}