//! The 12-byte DNS message header (RFC 1035 §4.1.1). Pure: no allocation,
//! no `std.Io`.
//!
//! Encode convention for the whole `dns/` module: fixed-size items encode into
//! a caller-provided buffer; variable-size items encode through a
//! `*std.Io.Writer`. The header is fixed-size, so it takes a `*[12]u8`.

const std = @import("std");
const types = @import("types.zig");

/// The second 16-bit word of the header, laid out from the least significant
/// bit up. `z` carries the three reserved bits (RFC 1035 §4.1.1; bits later
/// claimed as AD and CD by RFC 4035) so that an unknown bit pattern survives a
/// parse/encode round-trip unchanged.
pub const Flags = packed struct(u16) {
    rcode: types.Rcode,
    z: u3,
    ra: bool,
    rd: bool,
    tc: bool,
    aa: bool,
    opcode: types.Opcode,
    qr: bool,

    pub fn fromInt(value: u16) Flags {
        return @bitCast(value);
    }

    pub fn toInt(self: Flags) u16 {
        return @bitCast(self);
    }
};

pub const ParseError = error{Truncated};

pub const Header = struct {
    id: u16,
    flags: Flags,
    qdcount: u16,
    ancount: u16,
    nscount: u16,
    arcount: u16,
};

/// Reads the first `types.header_len` bytes. Trailing bytes are ignored.
pub fn parse(bytes: []const u8) ParseError!Header {
    if (bytes.len < types.header_len) return error.Truncated;
    const b = bytes[0..types.header_len];
    return .{
        .id = std.mem.readInt(u16, b[0..2], .big),
        .flags = Flags.fromInt(std.mem.readInt(u16, b[2..4], .big)),
        .qdcount = std.mem.readInt(u16, b[4..6], .big),
        .ancount = std.mem.readInt(u16, b[6..8], .big),
        .nscount = std.mem.readInt(u16, b[8..10], .big),
        .arcount = std.mem.readInt(u16, b[10..12], .big),
    };
}

pub fn encode(h: Header, out: *[types.header_len]u8) void {
    std.mem.writeInt(u16, out[0..2], h.id, .big);
    std.mem.writeInt(u16, out[2..4], h.flags.toInt(), .big);
    std.mem.writeInt(u16, out[4..6], h.qdcount, .big);
    std.mem.writeInt(u16, out[6..8], h.ancount, .big);
    std.mem.writeInt(u16, out[8..10], h.nscount, .big);
    std.mem.writeInt(u16, out[10..12], h.arcount, .big);
}

const testing = std.testing;

test "flags bit positions follow RFC 1035" {
    try testing.expectEqual(@as(u16, 0x8000), (Flags{
        .rcode = .no_error,
        .z = 0,
        .ra = false,
        .rd = false,
        .tc = false,
        .aa = false,
        .opcode = .query,
        .qr = true,
    }).toInt());

    try testing.expectEqual(@as(u16, 0x2800), (Flags{
        .rcode = .no_error,
        .z = 0,
        .ra = false,
        .rd = false,
        .tc = false,
        .aa = false,
        .opcode = .update,
        .qr = false,
    }).toInt());

    try testing.expectEqual(@as(u16, 0x0400), (Flags{
        .rcode = .no_error,
        .z = 0,
        .ra = false,
        .rd = false,
        .tc = false,
        .aa = true,
        .opcode = .query,
        .qr = false,
    }).toInt());

    try testing.expectEqual(@as(u16, 0x0200), (Flags{
        .rcode = .no_error,
        .z = 0,
        .ra = false,
        .rd = false,
        .tc = true,
        .aa = false,
        .opcode = .query,
        .qr = false,
    }).toInt());

    try testing.expectEqual(@as(u16, 0x0100), (Flags{
        .rcode = .no_error,
        .z = 0,
        .ra = false,
        .rd = true,
        .tc = false,
        .aa = false,
        .opcode = .query,
        .qr = false,
    }).toInt());

    try testing.expectEqual(@as(u16, 0x0080), (Flags{
        .rcode = .no_error,
        .z = 0,
        .ra = true,
        .rd = false,
        .tc = false,
        .aa = false,
        .opcode = .query,
        .qr = false,
    }).toInt());

    try testing.expectEqual(@as(u16, 0x0070), (Flags{
        .rcode = .no_error,
        .z = 7,
        .ra = false,
        .rd = false,
        .tc = false,
        .aa = false,
        .opcode = .query,
        .qr = false,
    }).toInt());

    try testing.expectEqual(@as(u16, 0x0003), (Flags{
        .rcode = .nx_domain,
        .z = 0,
        .ra = false,
        .rd = false,
        .tc = false,
        .aa = false,
        .opcode = .query,
        .qr = false,
    }).toInt());
}

test "every flags word round-trips, reserved bits included" {
    var value: u32 = 0;
    while (value <= std.math.maxInt(u16)) : (value += 1) {
        const word: u16 = @intCast(value);
        try testing.expectEqual(word, Flags.fromInt(word).toInt());
    }
}

test "header parse decodes a standard query" {
    const bytes = [_]u8{
        0xab, 0xcd, // id
        0x01, 0x20, // flags: RD set, z = 2
        0x00, 0x01, // qdcount
        0x00, 0x00, // ancount
        0x00, 0x00, // nscount
        0x00, 0x01, // arcount
    };
    const h = try parse(&bytes);
    try testing.expectEqual(@as(u16, 0xabcd), h.id);
    try testing.expectEqual(false, h.flags.qr);
    try testing.expectEqual(types.Opcode.query, h.flags.opcode);
    try testing.expectEqual(true, h.flags.rd);
    try testing.expectEqual(false, h.flags.ra);
    try testing.expectEqual(@as(u3, 2), h.flags.z);
    try testing.expectEqual(types.Rcode.no_error, h.flags.rcode);
    try testing.expectEqual(@as(u16, 1), h.qdcount);
    try testing.expectEqual(@as(u16, 0), h.ancount);
    try testing.expectEqual(@as(u16, 0), h.nscount);
    try testing.expectEqual(@as(u16, 1), h.arcount);
}

test "header round-trips through encode, unknown opcode and rcode included" {
    const original: Header = .{
        .id = 0x1234,
        .flags = .{
            .rcode = @enumFromInt(15),
            .z = 5,
            .ra = true,
            .rd = true,
            .tc = true,
            .aa = true,
            .opcode = @enumFromInt(3),
            .qr = true,
        },
        .qdcount = 1,
        .ancount = 2,
        .nscount = 3,
        .arcount = 4,
    };

    var buf: [types.header_len]u8 = undefined;
    encode(original, &buf);
    const parsed = try parse(&buf);
    try testing.expectEqual(original, parsed);
}

test "parse ignores bytes past the header" {
    const bytes = [_]u8{0xff} ** (types.header_len + 8);
    const h = try parse(&bytes);
    try testing.expectEqual(@as(u16, 0xffff), h.arcount);
}

test "truncated input" {
    var i: usize = 0;
    while (i < types.header_len) : (i += 1) {
        const bytes = [_]u8{0} ** types.header_len;
        try testing.expectError(error.Truncated, parse(bytes[0..i]));
    }
}

test "encode writes big-endian counts" {
    const h: Header = .{
        .id = 0x0102,
        .flags = Flags.fromInt(0x8180),
        .qdcount = 0x0304,
        .ancount = 0x0506,
        .nscount = 0x0708,
        .arcount = 0x090a,
    };
    var buf: [types.header_len]u8 = undefined;
    encode(h, &buf);
    try testing.expectEqualSlices(u8, &.{
        0x01, 0x02, 0x81, 0x80, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0a,
    }, &buf);
}
