package engine

import (
	"encoding/binary"
	"net"
	"testing"

	"github.com/google/gopacket/layers"
)

func TestParseL3IPv6UDP(t *testing.T) {
	src := net.ParseIP("2001:db8::10").To16()
	dst := net.ParseIP("2001:db8::20").To16()
	payload := []byte("hello")

	pkt := buildIPv6UDPPacket(t, src, dst, 12345, 443, payload)
	l3, transport, ok := ParseL3(pkt)
	if !ok {
		t.Fatal("ParseL3 should parse IPv6 packet")
	}
	if l3.Version != 6 {
		t.Fatalf("version=%d want=6", l3.Version)
	}
	if l3.Protocol != 17 {
		t.Fatalf("protocol=%d want=17 (udp)", l3.Protocol)
	}
	if !l3.SrcIPAddr().Equal(src) {
		t.Fatalf("src=%v want=%v", l3.SrcIPAddr(), src)
	}
	if !l3.DstIPAddr().Equal(dst) {
		t.Fatalf("dst=%v want=%v", l3.DstIPAddr(), dst)
	}
	udp, gotPayload, ok := ParseUDP(transport)
	if !ok {
		t.Fatal("ParseUDP should parse transport payload")
	}
	if udp.SrcPort != 12345 || udp.DstPort != 443 {
		t.Fatalf("ports=%d->%d want=12345->443", udp.SrcPort, udp.DstPort)
	}
	if string(gotPayload) != string(payload) {
		t.Fatalf("payload=%q want=%q", string(gotPayload), string(payload))
	}
}

func TestParseL3IPv6HopByHopThenUDP(t *testing.T) {
	src := net.ParseIP("2001:db8::1").To16()
	dst := net.ParseIP("2001:db8::2").To16()
	udpPayload := []byte("abc")
	udpLen := 8 + len(udpPayload)
	totalPayloadLen := 8 + udpLen // 8-byte hop-by-hop extension + udp packet

	pkt := make([]byte, 40+totalPayloadLen)
	pkt[0] = 0x60
	binary.BigEndian.PutUint16(pkt[4:6], uint16(totalPayloadLen))
	pkt[6] = 0  // Hop-by-hop
	pkt[7] = 64 // Hop limit
	copy(pkt[8:24], src)
	copy(pkt[24:40], dst)

	off := 40
	pkt[off+0] = 17 // Next header: UDP
	pkt[off+1] = 0  // Hdr Ext Len: (0+1)*8 = 8 bytes
	off += 8

	binary.BigEndian.PutUint16(pkt[off:off+2], 5353)
	binary.BigEndian.PutUint16(pkt[off+2:off+4], 53)
	binary.BigEndian.PutUint16(pkt[off+4:off+6], uint16(udpLen))
	copy(pkt[off+8:], udpPayload)

	l3, transport, ok := ParseL3(pkt)
	if !ok {
		t.Fatal("ParseL3 should parse IPv6 packet with extension headers")
	}
	if l3.Version != 6 || l3.Protocol != 17 {
		t.Fatalf("version/protocol=%d/%d want=6/17", l3.Version, l3.Protocol)
	}
	udp, gotPayload, ok := ParseUDP(transport)
	if !ok {
		t.Fatal("ParseUDP should parse UDP after hop-by-hop extension")
	}
	if udp.SrcPort != 5353 || udp.DstPort != 53 {
		t.Fatalf("ports=%d->%d want=5353->53", udp.SrcPort, udp.DstPort)
	}
	if string(gotPayload) != string(udpPayload) {
		t.Fatalf("payload=%q want=%q", string(gotPayload), string(udpPayload))
	}
}

func buildIPv6UDPPacket(t *testing.T, src, dst net.IP, srcPort, dstPort uint16, payload []byte) []byte {
	t.Helper()
	udpLen := 8 + len(payload)
	pkt := make([]byte, 40+udpLen)
	pkt[0] = 0x60
	binary.BigEndian.PutUint16(pkt[4:6], uint16(udpLen))
	pkt[6] = 17
	pkt[7] = 64
	copy(pkt[8:24], src.To16())
	copy(pkt[24:40], dst.To16())
	off := 40
	binary.BigEndian.PutUint16(pkt[off:off+2], srcPort)
	binary.BigEndian.PutUint16(pkt[off+2:off+4], dstPort)
	binary.BigEndian.PutUint16(pkt[off+4:off+6], uint16(udpLen))
	copy(pkt[off+8:], payload)
	return pkt
}

func TestParseL3IPv6FragmentNonFirst(t *testing.T) {
	src := net.ParseIP("2001:db8::a").To16()
	dst := net.ParseIP("2001:db8::b").To16()

	// IPv6 header + fragment header (offset != 0)
	pkt := make([]byte, 48)
	pkt[0] = 0x60
	binary.BigEndian.PutUint16(pkt[4:6], 8)
	pkt[6] = 44
	pkt[7] = 64
	copy(pkt[8:24], src)
	copy(pkt[24:40], dst)
	pkt[40] = uint8(layers.IPProtocolUDP)
	// fragment offset in 8-byte units: 1 (non-first fragment)
	binary.BigEndian.PutUint16(pkt[42:44], 1<<3)

	_, _, ok := ParseL3(pkt)
	if ok {
		t.Fatal("ParseL3 should reject non-first IPv6 fragments for L4 parsing")
	}
}