package udp

import (
	"encoding/binary"
	"testing"

	"git.difuse.io/Difuse/Mellaris/analyzer"
)

func makeWireGuardPacket(msgType byte, body []byte) []byte {
	buf := make([]byte, 1+3+len(body))
	buf[0] = msgType
	copy(buf[4:], body)
	return buf
}

func TestWireGuardUDPStream_Feed_HandshakeInitiation(t *testing.T) {
	s := newWireGuardUDPStream(nil)
	body := make([]byte, wireguardSizeHandshakeInitiation-4)
	binary.LittleEndian.PutUint32(body[0:4], 0x01020304)

	u, done := s.Feed(false, makeWireGuardPacket(wireguardTypeHandshakeInitiation, body))
	if u == nil {
		t.Fatal("Feed returned nil update")
	}
	if done {
		t.Error("should not be done")
	}
	if u.Type != analyzer.PropUpdateReplace {
		t.Errorf("Type = %d, want PropUpdateReplace", u.Type)
	}
	initMap, ok := u.M["handshake_initiation"].(analyzer.PropMap)
	if !ok {
		t.Fatal("handshake_initiation missing")
	}
	if idx := initMap["sender_index"].(uint32); idx != 0x01020304 {
		t.Errorf("sender_index = %d, want %d", idx, 0x01020304)
	}
}

func TestWireGuardUDPStream_Feed_HandshakeResponse(t *testing.T) {
	s := newWireGuardUDPStream(nil)
	body := make([]byte, wireguardSizeHandshakeResponse-4)
	binary.LittleEndian.PutUint32(body[0:4], 0x01020304)
	binary.LittleEndian.PutUint32(body[4:8], 0x05060708)

	u, _ := s.Feed(false, makeWireGuardPacket(wireguardTypeHandshakeResponse, body))
	if u == nil {
		t.Fatal("Feed returned nil update")
	}
	respMap, ok := u.M["handshake_response"].(analyzer.PropMap)
	if !ok {
		t.Fatal("handshake_response missing")
	}
	if idx := respMap["sender_index"].(uint32); idx != 0x01020304 {
		t.Errorf("sender_index = %d", idx)
	}
	if idx := respMap["receiver_index"].(uint32); idx != 0x05060708 {
		t.Errorf("receiver_index = %d", idx)
	}
}

func TestWireGuardUDPStream_Feed_PacketData(t *testing.T) {
	s := newWireGuardUDPStream(nil)
	body := make([]byte, wireguardMinSizePacketData-4)
	binary.LittleEndian.PutUint32(body[0:4], 0x0a0b0c0d)
	binary.LittleEndian.PutUint64(body[4:12], 42)

	u, _ := s.Feed(false, makeWireGuardPacket(wireguardTypeData, body))
	if u == nil {
		t.Fatal("Feed returned nil update")
	}
	dataMap, ok := u.M["packet_data"].(analyzer.PropMap)
	if !ok {
		t.Fatal("packet_data missing")
	}
	if idx := dataMap["receiver_index"].(uint32); idx != 0x0a0b0c0d {
		t.Errorf("receiver_index = %d", idx)
	}
	if ctr := dataMap["counter"].(uint64); ctr != 42 {
		t.Errorf("counter = %d", ctr)
	}
}

func TestWireGuardUDPStream_Feed_CookieReply(t *testing.T) {
	s := newWireGuardUDPStream(nil)
	body := make([]byte, wireguardSizePacketCookieReply-4)
	binary.LittleEndian.PutUint32(body[0:4], 0x11111111)

	u, _ := s.Feed(false, makeWireGuardPacket(wireguardTypeCookieReply, body))
	if u == nil {
		t.Fatal("Feed returned nil update")
	}
	crMap, ok := u.M["packet_cookie_reply"].(analyzer.PropMap)
	if !ok {
		t.Fatal("packet_cookie_reply missing")
	}
	if idx := crMap["receiver_index"].(uint32); idx != 0x11111111 {
		t.Errorf("receiver_index = %d", idx)
	}
}

func TestWireGuardUDPStream_Feed_InvalidLength(t *testing.T) {
	s := newWireGuardUDPStream(nil)
	u, done := s.Feed(false, []byte{wireguardTypeHandshakeInitiation, 0, 0, 0, 0x01})
	if u == nil || !done {
		u, done = s.Feed(false, []byte{wireguardTypeHandshakeInitiation, 0, 0, 0, 0x01})
	}
	if u == nil || !done {
		u, _ = s.Feed(false, []byte{wireguardTypeHandshakeInitiation, 0, 0, 0, 0x01})
	}
	u, done = s.Feed(false, []byte{wireguardTypeHandshakeInitiation, 0, 0, 0, 0x01})
	if u != nil || !done {
		t.Error("should return done after 4 invalid packets")
	}
}

func TestWireGuardUDPStream_Feed_TooShort(t *testing.T) {
	s := newWireGuardUDPStream(nil)
	u, _ := s.Feed(false, []byte{1, 0, 0})
	if u != nil {
		t.Error("should return nil for too-short packet")
	}
}

func TestWireGuardUDPStream_Feed_NonZeroReserved(t *testing.T) {
	s := newWireGuardUDPStream(nil)
	u, _ := s.Feed(false, []byte{wireguardTypeHandshakeInitiation, 0, 0, 1, 0, 0, 0, 0})
	if u != nil {
		t.Error("should return nil when reserved bytes are non-zero")
	}
}

func TestWireGuardUDPStream_Feed_HandshakeInitiation_WrongSize(t *testing.T) {
	s := newWireGuardUDPStream(nil)
	body := make([]byte, wireguardSizeHandshakeInitiation-4+1)
	binary.LittleEndian.PutUint32(body[0:4], 1)
	u, _ := s.Feed(false, makeWireGuardPacket(wireguardTypeHandshakeInitiation, body))
	if u != nil {
		t.Error("should return nil for wrong init size")
	}
}

func TestWireGuardUDPStream_Close(t *testing.T) {
	s := newWireGuardUDPStream(nil)
	u := s.Close(false)
	if u != nil {
		t.Error("Close() should return nil")
	}
}

func TestWireGuardUDPStream_PutSenderIndex_MatchReceiverIndex(t *testing.T) {
	s := newWireGuardUDPStream(nil)

	s.putSenderIndex(false, 0xdeadbeef)
	if !s.matchReceiverIndex(true, 0xdeadbeef) {
		t.Error("should match reverse direction")
	}
	if s.matchReceiverIndex(false, 0xdeadbeef) {
		t.Error("should not match same direction")
	}
	if s.matchReceiverIndex(true, 0x12345678) {
		t.Error("should not match wrong index")
	}
}

func TestWireGuardUDPStream_RememberedIndexRing(t *testing.T) {
	s := newWireGuardUDPStream(nil)

	for i := uint32(0); i < wireguardRememberedIndexCount+2; i++ {
		s.putSenderIndex(false, i)
	}

	if s.matchReceiverIndex(true, 0) {
		t.Error("index 0 should have been evicted from ring")
	}
	if !s.matchReceiverIndex(true, uint32(wireguardRememberedIndexCount+1)) {
		t.Error("latest index should still be in ring")
	}
}

func TestWireGuardAnalyzer_Name(t *testing.T) {
	a := &WireGuardAnalyzer{}
	if a.Name() != "wireguard" {
		t.Errorf("Name() = %q, want wireguard", a.Name())
	}
}