// SPDX-FileCopyrightText: 2026 Echolot contributors
// SPDX-License-Identifier: GPL-3.0-or-later

package canarydns

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

// buildQuery makes a single-question DNS query, optionally with an EDNS OPT.
func buildQuery(name string, qtype uint16, ednsBufsize int) []byte {
	msg := make([]byte, 12)
	binary.BigEndian.PutUint16(msg[0:2], 0x1234)
	binary.BigEndian.PutUint16(msg[2:4], flagRD)
	binary.BigEndian.PutUint16(msg[4:6], 1) // QDCOUNT
	msg = append(msg, encodeName(name)...)
	msg = binary.BigEndian.AppendUint16(msg, qtype)
	msg = binary.BigEndian.AppendUint16(msg, classIN)
	if ednsBufsize > 0 {
		binary.BigEndian.PutUint16(msg[10:12], 1) // ARCOUNT
		msg = append(msg, 0)                      // root name
		msg = binary.BigEndian.AppendUint16(msg, typeOPT)
		msg = binary.BigEndian.AppendUint16(msg, uint16(ednsBufsize))
		msg = binary.BigEndian.AppendUint32(msg, 0)
		msg = binary.BigEndian.AppendUint16(msg, 0)
	}
	return msg
}

// parseAnswers pulls (type, ttl, rdata) tuples from a response.
type ans struct {
	typ  uint16
	ttl  uint32
	data []byte
}

func parseResponse(t *testing.T, resp []byte) (flags uint16, answers []ans) {
	t.Helper()
	flags = binary.BigEndian.Uint16(resp[2:4])
	qd := binary.BigEndian.Uint16(resp[4:6])
	an := binary.BigEndian.Uint16(resp[6:8])
	off := 12
	for i := uint16(0); i < qd; i++ {
		_, next, ok := readName(resp, off)
		if !ok {
			t.Fatal("bad question name")
		}
		off = next + 4
	}
	for i := uint16(0); i < an; i++ {
		_, next, ok := readName(resp, off)
		if !ok {
			t.Fatal("bad answer name")
		}
		typ := binary.BigEndian.Uint16(resp[next : next+2])
		ttl := binary.BigEndian.Uint32(resp[next+4 : next+8])
		rdlen := int(binary.BigEndian.Uint16(resp[next+8 : next+10]))
		answers = append(answers, ans{typ, ttl, resp[next+10 : next+10+rdlen]})
		off = next + 10 + rdlen
	}
	return
}

func newTestServer() *Server {
	return New("c.echo-lot.app", "fmr", netip.MustParseAddr("192.0.2.1"), netip.MustParseAddr("2001:db8::1"), 24*time.Hour)
}

func TestReferenceRecords(t *testing.T) {
	s := newTestServer()
	resolver := netip.MustParseAddr("198.51.100.7")

	// ttl-5 A → 192.0.2.5, TTL 5
	resp := s.handle(buildQuery("ttl-5.c.echo-lot.app", typeA, 0), resolver, "udp")
	_, answers := parseResponse(t, resp)
	if len(answers) != 1 || answers[0].ttl != 5 || !netip.AddrFrom4([4]byte(answers[0].data)).IsValid() {
		t.Fatalf("ttl-5 A: %+v", answers)
	}
	if got := net.IP(answers[0].data).String(); got != "192.0.2.5" {
		t.Fatalf("ttl-5 A = %s, want 192.0.2.5", got)
	}

	// many-rr → exactly 8 A records, in order .101..108
	resp = s.handle(buildQuery("many-rr.c.echo-lot.app", typeA, 0), resolver, "udp")
	_, answers = parseResponse(t, resp)
	if len(answers) != 8 {
		t.Fatalf("many-rr: got %d A records, want 8", len(answers))
	}
	for i, a := range answers {
		if a.data[3] != byte(101+i) {
			t.Fatalf("many-rr order: record %d = .%d, want .%d", i, a.data[3], 101+i)
		}
	}
}

func TestBigTxtTruncationVsEDNS(t *testing.T) {
	s := newTestServer()
	resolver := netip.MustParseAddr("198.51.100.7")

	// No EDNS → 512 cap → TC set, answers dropped.
	resp := s.handle(buildQuery("big-txt.c.echo-lot.app", typeTXT, 0), resolver, "udp")
	flags, answers := parseResponse(t, resp)
	if flags&flagTC == 0 {
		t.Fatal("big-txt over plain UDP should set TC")
	}
	if len(answers) != 0 {
		t.Fatalf("truncated response should carry no answers, got %d", len(answers))
	}

	// EDNS bufsize 4096 → full answer, no TC.
	resp = s.handle(buildQuery("big-txt.c.echo-lot.app", typeTXT, 4096), resolver, "udp")
	flags, answers = parseResponse(t, resp)
	if flags&flagTC != 0 {
		t.Fatal("big-txt with EDNS 4096 should not truncate")
	}
	if len(answers) != 1 {
		t.Fatalf("want 1 TXT answer, got %d", len(answers))
	}

	// TCP → never truncates.
	resp = s.handle(buildQuery("big-txt.c.echo-lot.app", typeTXT, 0), resolver, "tcp")
	flags, _ = parseResponse(t, resp)
	if flags&flagTC != 0 {
		t.Fatal("TCP must not truncate")
	}
}

func TestQueryLogAndPerPrefix(t *testing.T) {
	s := newTestServer()
	s.handle(buildQuery("abc123.SESSPREFIX1.c.echo-lot.app", typeA, 1232), netip.MustParseAddr("198.51.100.7"), "udp")
	s.handle(buildQuery("def456.other.c.echo-lot.app", typeA, 0), netip.MustParseAddr("203.0.113.9"), "udp")

	all := s.RecentForPrefix("sessprefix1")
	if len(all) != 1 {
		t.Fatalf("per-prefix filter: got %d, want 1", len(all))
	}
	q := all[0]
	if q.ResolverIP != "198.51.100.7" || q.Transport != "udp" {
		t.Fatalf("logged resolver/transport wrong: %+v", q)
	}
	if !q.EDNS.Present || q.EDNS.Bufsize != 1232 {
		t.Fatalf("EDNS not captured: %+v", q.EDNS)
	}
	// nonce answer is deterministic + in doc range
	resp := s.handle(buildQuery("abc123.SESSPREFIX1.c.echo-lot.app", typeA, 0), netip.MustParseAddr("198.51.100.7"), "udp")
	_, answers := parseResponse(t, resp)
	if len(answers) != 1 || answers[0].data[0] != 192 || answers[0].data[1] != 0 || answers[0].data[2] != 2 {
		t.Fatalf("nonce answer not in 192.0.2.0/24: %+v", answers)
	}
}

func TestOutOfZoneNXDomain(t *testing.T) {
	s := newTestServer()
	resp := s.handle(buildQuery("example.com", typeA, 0), netip.MustParseAddr("198.51.100.7"), "udp")
	flags, _ := parseResponse(t, resp)
	if flags&0x000F != rcodeNXDomain {
		t.Fatalf("out-of-zone should be NXDOMAIN, flags=%#x", flags)
	}
}

func TestQueryLogRetentionForgetsOldEntries(t *testing.T) {
	s := newTestServer() // 24 h retention
	now := time.Now().UTC()
	s.record(Query{QName: "old.sess1.c.echo-lot.app", At: now.Add(-25 * time.Hour)})
	s.record(Query{QName: "fresh.sess1.c.echo-lot.app", At: now})

	got := s.RecentForPrefix("sess1")
	if len(got) != 1 || got[0].QName != "fresh.sess1.c.echo-lot.app" {
		t.Fatalf("retention not enforced: %+v", got)
	}

	// Reads must age the log too: an idle server still has to forget on schedule.
	s.log[0].At = now.Add(-25 * time.Hour)
	if got := s.RecentForPrefix("sess1"); len(got) != 0 {
		t.Fatalf("read path did not expire entries: %+v", got)
	}
}

func TestZeroRetentionKeepsEverything(t *testing.T) {
	s := New("c.echo-lot.app", "fmr", netip.MustParseAddr("192.0.2.1"), netip.MustParseAddr("2001:db8::1"), 0)
	s.record(Query{QName: "ancient.sess1.c.echo-lot.app", At: time.Now().UTC().Add(-1000 * time.Hour)})
	if got := s.RecentForPrefix("sess1"); len(got) != 1 {
		t.Fatal("retention 0 must mean 'ring cap only', not 'keep nothing'")
	}
}