package dnstt

import (
	"encoding/binary"
	"net/netip"
	"testing"
	"time"
)

func TestProcessIPv6PacketRecordsResolverAddress(t *testing.T) {
	now := time.Unix(1000, 0)
	resolverIP := netip.MustParseAddr("2001:db8::53")
	serverIP := netip.MustParseAddr("2001:db8::1")
	c := NewCollector(
		[]string{"tunnel.example.com"},
		WithNow(func() time.Time { return now }),
		WithGeoResolver(fakeGeoResolver{
			labelNames: []string{"asn"},
			labels: map[netip.Addr]GeoLabels{
				resolverIP: {ASN: "13335"},
			},
		}),
	)

	ProcessPacket(ipv6UDPPacket(resolverIP, serverIP, 53000, 53, dnsQuery("MFRGGZDFMZTWQ2LK", "tunnel.example.com")), 53, c)

	series := c.Snapshot().Tunnels["tunnel.example.com"].Series
	if len(series) != 1 {
		t.Fatalf("series count = %d, want 1", len(series))
	}
	if series[0].ASN != "13335" {
		t.Fatalf("ASN = %q, want 13335", series[0].ASN)
	}
	if series[0].TotalQueries != 1 {
		t.Fatalf("queries = %d, want 1", series[0].TotalQueries)
	}
}

func ipv6UDPPacket(src, dst netip.Addr, srcPort, dstPort uint16, payload []byte) []byte {
	packet := make([]byte, 40+8+len(payload))
	packet[0] = 0x60
	binary.BigEndian.PutUint16(packet[4:6], uint16(8+len(payload)))
	packet[6] = 17
	packet[7] = 64
	srcBytes := src.As16()
	dstBytes := dst.As16()
	copy(packet[8:24], srcBytes[:])
	copy(packet[24:40], dstBytes[:])

	udp := packet[40:]
	binary.BigEndian.PutUint16(udp[0:2], srcPort)
	binary.BigEndian.PutUint16(udp[2:4], dstPort)
	binary.BigEndian.PutUint16(udp[4:6], uint16(8+len(payload)))
	copy(udp[8:], payload)
	return packet
}