-- SPDX-FileCopyrightText: © 2022—2026 Vladimir Zorin <vladimir@deviant.guru>
-- SPDX-License-Identifier: LicenseRef-OWL-1.0-or-later
-- Licensed under OWL v1.0+. See LICENSE.

--[==[
Full DNS message encode and decode. Handles headers, question sections,
resource records with RDATA codec dispatch, and EDNS OPT pseudo-record
extraction/injection.
]==]

local wire = require("dns.wire")
local types = require("dns.types")

local bit = require("bit")
local band, bor, lshift, rshift = bit.band, bit.bor, bit.lshift, bit.rshift

local TYPE = types.TYPE

---------------------------------------------------------------------------
-- Internal: header decode/encode
---------------------------------------------------------------------------

---! Decode the 12-byte DNS header from a reader
---@ decode_header(`r`{.tbl}) -> `hdr`{.tbl .or_nil}, `err`{.str .err}
local decode_header = function(r)
	local id, err = r:u16()
	if not id then
		return nil, "header: " .. err
	end
	local flags
	flags, err = r:u16()
	if not flags then
		return nil, "header: " .. err
	end
	local qdcount
	qdcount, err = r:u16()
	if not qdcount then
		return nil, "header: " .. err
	end
	local ancount
	ancount, err = r:u16()
	if not ancount then
		return nil, "header: " .. err
	end
	local nscount
	nscount, err = r:u16()
	if not nscount then
		return nil, "header: " .. err
	end
	local arcount
	arcount, err = r:u16()
	if not arcount then
		return nil, "header: " .. err
	end

	return {
		id = id,
		qr = band(rshift(flags, 15), 0x1),
		opcode = band(rshift(flags, 11), 0xF),
		aa = band(rshift(flags, 10), 0x1) == 1,
		tc = band(rshift(flags, 9), 0x1) == 1,
		rd = band(rshift(flags, 8), 0x1) == 1,
		ra = band(rshift(flags, 7), 0x1) == 1,
		rcode = band(flags, 0xF),
	},
		qdcount,
		ancount,
		nscount,
		arcount
end

---! Encode a DNS header into a writer
---@ encode_header(`w`{.tbl}, `hdr`{.tbl}, `qdcount`{.num}, `ancount`{.num}, `nscount`{.num}, `arcount`{.num})
local encode_header = function(w, hdr, qdcount, ancount, nscount, arcount)
	w:u16(hdr.id)
	w:u16(
		bor(
			lshift(hdr.qr or 0, 15),
			lshift(hdr.opcode or 0, 11),
			lshift(hdr.aa and 1 or 0, 10),
			lshift(hdr.tc and 1 or 0, 9),
			lshift(hdr.rd and 1 or 0, 8),
			lshift(hdr.ra and 1 or 0, 7),
			band(hdr.rcode or 0, 0xF)
		)
	)
	w:u16(qdcount)
	w:u16(ancount)
	w:u16(nscount)
	w:u16(arcount)
end

---------------------------------------------------------------------------
-- Internal: question decode/encode
---------------------------------------------------------------------------

---! Decode a single question entry
---@ decode_question(`r`{.tbl}) -> `q`{.tbl .or_nil}, `err`{.str .err}
local decode_question = function(r)
	local name, err = r:name()
	if not name then
		return nil, "question: " .. err
	end
	local qtype
	qtype, err = r:u16()
	if not qtype then
		return nil, "question: " .. err
	end
	local qclass
	qclass, err = r:u16()
	if not qclass then
		return nil, "question: " .. err
	end
	return { name = name, qtype = qtype, qclass = qclass }
end

---! Encode a single question entry
---@ encode_question(`w`{.tbl}, `q`{.tbl})
local encode_question = function(w, q)
	w:name(q.name)
	w:u16(q.qtype)
	w:u16(q.qclass)
end

---------------------------------------------------------------------------
-- Internal: resource record decode/encode
---------------------------------------------------------------------------

---! Decode a single resource record
---@ decode_rr(`r`{.tbl}) -> `rr`{.tbl .or_nil}, `err`{.str .err}
local decode_rr = function(r)
	local name, err = r:name()
	if not name then
		return nil, "rr: " .. err
	end
	local rtype
	rtype, err = r:u16()
	if not rtype then
		return nil, "rr: " .. err
	end
	local rclass
	rclass, err = r:u16()
	if not rclass then
		return nil, "rr: " .. err
	end
	local ttl
	ttl, err = r:u32()
	if not ttl then
		return nil, "rr: " .. err
	end
	local rdlength
	rdlength, err = r:u16()
	if not rdlength then
		return nil, "rr: " .. err
	end

	local rdata
	local codec = types.get_codec(rtype)
	if codec and rdlength > 0 then
		local sub_r
		sub_r, err = r:sub_reader(rdlength)
		if not sub_r then
			return nil, "rr: " .. err
		end
		rdata, err = codec.decode(sub_r)
		if not rdata then
			return nil, "rr rdata: " .. err
		end
	elseif rdlength > 0 then
		-- Unknown type: store raw bytes
		local raw_bytes
		raw_bytes, err = r:bytes(rdlength)
		if not raw_bytes then
			return nil, "rr: " .. err
		end
		rdata = { raw = raw_bytes }
	else
		-- Zero-length RDATA
		if codec then
			-- Let codec handle empty (e.g. OPT with no options)
			local sub_r
			sub_r, err = r:sub_reader(0)
			if not sub_r then
				return nil, "rr: " .. err
			end
			rdata, err = codec.decode(sub_r)
			if not rdata then
				return nil, "rr rdata: " .. err
			end
		else
			rdata = { raw = "" }
		end
	end

	return {
		name = name,
		type = rtype,
		class = rclass,
		ttl = ttl,
		rdata = rdata,
	}
end

---! Encode a single resource record
---@ encode_rr(`w`{.tbl}, `rr`{.tbl})
local encode_rr = function(w, rr)
	w:name(rr.name)
	w:u16(rr.type)
	w:u16(rr.class)
	w:u32(rr.ttl)

	-- Write RDLENGTH placeholder, then RDATA, then patch
	local rdlen_offset = w:pos()
	w:u16(0) -- placeholder

	local codec = types.get_codec(rr.type)
	if codec then
		codec.encode(w, rr.rdata)
	elseif rr.rdata and rr.rdata.raw then
		w:bytes(rr.rdata.raw)
	end

	local rdlength = w:pos() - rdlen_offset - 2
	w:patch_u16(rdlen_offset, rdlength)
end

---------------------------------------------------------------------------
-- Internal: EDNS OPT extraction/injection
---------------------------------------------------------------------------

---! Extract OPT pseudo-record from additional section into edns table
---@ extract_edns(`additional`{.tbl}) -> `edns`{.tbl .or_nil}, `filtered`{.tbl}
local extract_edns = function(additional)
	local edns = nil
	local filtered = {}
	for _, rr in ipairs(additional) do
		if rr.type == TYPE.OPT then
			-- CLASS → udp_size
			-- TTL bits: 24-31 extended_rcode, 16-23 version, 15 DO, 0-14 reserved
			local ttl_val = rr.ttl
			edns = {
				udp_size = rr.class,
				extended_rcode = band(rshift(ttl_val, 24), 0xFF),
				version = band(rshift(ttl_val, 16), 0xFF),
				do_bit = band(rshift(ttl_val, 15), 0x1) == 1,
				options = rr.rdata and rr.rdata.options or {},
			}
		else
			filtered[#filtered + 1] = rr
		end
	end
	return edns, filtered
end

---! Build OPT pseudo-record RR from edns table
---@ build_opt_rr(`edns`{.tbl}) -> `rr`{.tbl}
local build_opt_rr = function(edns)
	local ttl_val = bor(
		lshift(band(edns.extended_rcode or 0, 0xFF), 24),
		lshift(band(edns.version or 0, 0xFF), 16),
		lshift(edns.do_bit and 1 or 0, 15)
	)
	return {
		name = ".",
		type = TYPE.OPT,
		class = edns.udp_size or 4096,
		ttl = ttl_val,
		rdata = { options = edns.options or {} },
	}
end

---------------------------------------------------------------------------
-- Decode
---------------------------------------------------------------------------

---! Decode a raw DNS message into a structured table
---@ decode(`raw_bytes`{.str}) -> `msg`{.tbl .or_nil}, `err`{.str .err}
--[===[
Decodes a complete DNS message from a raw binary string. Returns a table
with `header`, `question`, `answer`, `authority`, `additional`, and
optionally `edns` fields. The OPT pseudo-record is extracted from the
additional section into `msg.edns`.

Section counts from the header determine how many entries are parsed in
each section. If any entry fails to parse, the entire decode fails with
an error (strict mode).
]===]
local decode = function(raw_bytes)
	if not raw_bytes or #raw_bytes < 12 then
		return nil, "message too short"
	end

	local r = wire.reader(raw_bytes)

	-- Header
	local hdr, err_or_qdcount, ancount, nscount, arcount = decode_header(r)
	if not hdr then
		return nil, err_or_qdcount
	end
	local qdcount = err_or_qdcount

	-- Question section
	local question = {}
	for _ = 1, qdcount do
		local q, err = decode_question(r)
		if not q then
			return nil, err
		end
		question[#question + 1] = q
	end

	-- Decode RR sections
	local decode_section = function(count)
		local section = {}
		for _ = 1, count do
			local rr, err = decode_rr(r)
			if not rr then
				return nil, err
			end
			section[#section + 1] = rr
		end
		return section
	end

	local answer, err = decode_section(ancount)
	if not answer then
		return nil, err
	end

	local authority
	authority, err = decode_section(nscount)
	if not authority then
		return nil, err
	end

	local additional
	additional, err = decode_section(arcount)
	if not additional then
		return nil, err
	end

	-- Extract EDNS OPT
	local edns
	edns, additional = extract_edns(additional)

	return {
		header = hdr,
		question = question,
		answer = answer,
		authority = authority,
		additional = additional,
		edns = edns,
	}
end

---------------------------------------------------------------------------
-- Encode
---------------------------------------------------------------------------

---! Encode a structured message table into raw DNS wire format
---@ encode(`msg`{.tbl}) -> `raw_bytes`{.str}
--[===[
Encodes a complete DNS message from a structured table. Section counts
are computed from the array lengths. If `msg.edns` is set, an OPT
pseudo-record is appended to the additional section.
]===]
local encode = function(msg)
	local w = wire.writer()

	local question = msg.question or {}
	local answer = msg.answer or {}
	local authority = msg.authority or {}
	local additional = msg.additional or {}

	local arcount = #additional
	if msg.edns then
		arcount = arcount + 1
	end

	-- Header
	encode_header(w, msg.header, #question, #answer, #authority, arcount)

	-- Question section
	for _, q in ipairs(question) do
		encode_question(w, q)
	end

	-- Answer section
	for _, rr in ipairs(answer) do
		encode_rr(w, rr)
	end

	-- Authority section
	for _, rr in ipairs(authority) do
		encode_rr(w, rr)
	end

	-- Additional section
	for _, rr in ipairs(additional) do
		encode_rr(w, rr)
	end

	-- EDNS OPT
	if msg.edns then
		encode_rr(w, build_opt_rr(msg.edns))
	end

	return w:output()
end

---------------------------------------------------------------------------
-- Module Export
---------------------------------------------------------------------------

return {
	decode = decode,
	encode = encode,
}