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

package canarydns

import (
	"encoding/binary"
	"net/netip"
	"strings"
)

// rr is a resource record to encode into the answer section.
type rr struct {
	ttl  uint32
	typ  uint16
	addr netip.Addr // for A/AAAA
	txt  []string   // for TXT
	ns   string     // for NS
}

// optInfo is the parsed EDNS OPT plus the derived observation fields.
type optInfo struct {
	edns edns
	ecs  string
}

// parseQuestion reads a single question starting at off. Returns the raw
// (case-preserved) qname with trailing dot, qtype, qclass, and the offset
// just past the question.
func parseQuestion(pkt []byte, off int) (qname string, qtype, qclass uint16, end int, ok bool) {
	name, next, ok := readName(pkt, off)
	if !ok || next+4 > len(pkt) {
		return "", 0, 0, 0, false
	}
	qtype = binary.BigEndian.Uint16(pkt[next : next+2])
	qclass = binary.BigEndian.Uint16(pkt[next+2 : next+4])
	return name, qtype, qclass, next + 4, true
}

// readName decodes a DNS name (with compression pointers) into a
// dot-terminated string, preserving label case.
func readName(pkt []byte, off int) (string, int, bool) {
	var sb strings.Builder
	end := -1
	jumps := 0
	for {
		if off >= len(pkt) {
			return "", 0, false
		}
		l := int(pkt[off])
		switch {
		case l == 0:
			off++
			if end < 0 {
				end = off
			}
			if sb.Len() == 0 {
				return ".", end, true
			}
			return sb.String(), end, true
		case l&0xC0 == 0xC0: // compression pointer
			if off+1 >= len(pkt) {
				return "", 0, false
			}
			if end < 0 {
				end = off + 2
			}
			off = int(binary.BigEndian.Uint16(pkt[off:off+2]) & 0x3FFF)
			jumps++
			if jumps > 16 {
				return "", 0, false
			}
		default:
			if off+1+l > len(pkt) {
				return "", 0, false
			}
			sb.Write(pkt[off+1 : off+1+l])
			sb.WriteByte('.')
			off += 1 + l
		}
	}
}

// parseOPT scans the additional section for an EDNS OPT RR and extracts
// bufsize, the DO flag, and any ECS option.
func parseOPT(pkt []byte, off int, arcount uint16) optInfo {
	var info optInfo
	for i := uint16(0); i < arcount && off < len(pkt); i++ {
		_, next, ok := readName(pkt, off)
		if !ok || next+10 > len(pkt) {
			return info
		}
		typ := binary.BigEndian.Uint16(pkt[next : next+2])
		class := binary.BigEndian.Uint16(pkt[next+2 : next+4]) // OPT: requester bufsize
		ttl := binary.BigEndian.Uint32(pkt[next+4 : next+8])   // OPT: extended-rcode/version/flags
		rdlen := int(binary.BigEndian.Uint16(pkt[next+8 : next+10]))
		rdata := next + 10
		if rdata+rdlen > len(pkt) {
			return info
		}
		if typ == typeOPT {
			info.edns.Present = true
			info.edns.Bufsize = int(class)
			if ttl&ednsDO != 0 {
				info.edns.Flags = append(info.edns.Flags, "do")
			}
			info.ecs = parseECS(pkt[rdata : rdata+rdlen])
			return info
		}
		off = rdata + rdlen
	}
	return info
}

// parseECS extracts an EDNS Client Subnet option (RFC 7871) as "ip/scope".
func parseECS(rdata []byte) string {
	for len(rdata) >= 4 {
		code := binary.BigEndian.Uint16(rdata[0:2])
		olen := int(binary.BigEndian.Uint16(rdata[2:4]))
		if 4+olen > len(rdata) {
			return ""
		}
		if code == optECS && olen >= 4 {
			fam := binary.BigEndian.Uint16(rdata[4:6])
			srcPrefix := rdata[6]
			addrBytes := rdata[8 : 4+olen]
			var ip netip.Addr
			if fam == 1 {
				var b [4]byte
				copy(b[:], addrBytes)
				ip = netip.AddrFrom4(b)
			} else if fam == 2 {
				var b [16]byte
				copy(b[:], addrBytes)
				ip = netip.AddrFrom16(b)
			}
			if ip.IsValid() {
				return ip.String() + "/" + itoa(int(srcPrefix))
			}
		}
		rdata = rdata[4+olen:]
	}
	return ""
}

func itoa(n int) string {
	if n == 0 {
		return "0"
	}
	var b [4]byte
	i := len(b)
	for n > 0 {
		i--
		b[i] = byte('0' + n%10)
		n /= 10
	}
	return string(b[i:])
}

// buildResponse assembles the answer, sets TC when a UDP response exceeds the
// negotiated buffer, and appends the OPT RR when the query used EDNS.
func (s *Server) buildResponse(id uint16, pkt []byte, qEnd int, answers []rr, opt *optInfo, transport string, rcode int) []byte {
	msg := make([]byte, 12)
	binary.BigEndian.PutUint16(msg[0:2], id)
	// question is copied verbatim (case preserved) from the query
	msg = append(msg, pkt[12:qEnd]...)

	body := make([]byte, 0, 512)
	for _, a := range answers {
		body = append(body, encodeRR(a)...)
	}

	extra := 0
	if opt != nil && opt.edns.Present {
		extra = 1
	}

	flags := uint16(flagQR|flagAA) | (binary.BigEndian.Uint16(pkt[2:4]) & flagRD) | uint16(rcode)
	if opt != nil && opt.edns.Present {
		body = append(body, buildOPT(opt)...)
	}

	// UDP truncation: without EDNS the limit is 512; with EDNS it's the
	// requester's bufsize (floored at 512). Drop the answer section and set TC.
	if transport == "udp" {
		limit := udpMaxNoEDNS
		if opt != nil && opt.edns.Present && opt.edns.Bufsize > udpMaxNoEDNS {
			limit = opt.edns.Bufsize
		}
		if 12+(qEnd-12)+len(body) > limit {
			flags |= flagTC
			// Keep only the OPT RR (if any); drop answers.
			body = body[:0]
			if opt != nil && opt.edns.Present {
				body = append(body, buildOPT(opt)...)
				answers = nil
			} else {
				answers = nil
			}
		}
	}

	binary.BigEndian.PutUint16(msg[2:4], uint16(flags))
	binary.BigEndian.PutUint16(msg[4:6], 1) // QDCOUNT
	binary.BigEndian.PutUint16(msg[6:8], uint16(len(answers)))
	binary.BigEndian.PutUint16(msg[10:12], uint16(extra))
	return append(msg, body...)
}

// encodeRR encodes one answer RR, using a compression pointer (0xC00C) to the
// question name at offset 12.
func encodeRR(a rr) []byte {
	var rdata []byte
	switch a.typ {
	case typeA:
		b := a.addr.As4()
		rdata = b[:]
	case typeAAAA:
		b := a.addr.As16()
		rdata = b[:]
	case typeTXT:
		for _, s := range a.txt {
			for len(s) > 0 {
				n := len(s)
				if n > 255 {
					n = 255
				}
				rdata = append(rdata, byte(n))
				rdata = append(rdata, s[:n]...)
				s = s[n:]
			}
		}
	case typeNS:
		rdata = encodeName(a.ns)
	}
	out := make([]byte, 0, 12+len(rdata))
	out = append(out, 0xC0, 0x0C) // name → pointer to question
	out = binary.BigEndian.AppendUint16(out, a.typ)
	out = binary.BigEndian.AppendUint16(out, classIN)
	out = binary.BigEndian.AppendUint32(out, a.ttl)
	out = binary.BigEndian.AppendUint16(out, uint16(len(rdata)))
	return append(out, rdata...)
}

func encodeName(name string) []byte {
	var out []byte
	for _, label := range strings.Split(strings.TrimSuffix(name, "."), ".") {
		if label == "" {
			continue
		}
		out = append(out, byte(len(label)))
		out = append(out, label...)
	}
	return append(out, 0)
}

// buildOPT emits a minimal EDNS OPT RR echoing our own bufsize (advertise a
// generous 4096) with DO cleared — we serve no DNSSEC.
func buildOPT(*optInfo) []byte {
	out := []byte{0} // root name
	out = binary.BigEndian.AppendUint16(out, typeOPT)
	out = binary.BigEndian.AppendUint16(out, 4096) // our bufsize
	out = binary.BigEndian.AppendUint32(out, 0)    // ext-rcode/version/flags
	out = binary.BigEndian.AppendUint16(out, 0)    // rdlen
	return out
}