//! Local DNS records (PLAN §6.4): an immutable, sorted lookup table built once
//! from the `local_records` rows and read on the query path without allocating.
//! Pure: an allocator and plain values, no `std.Io`, no clock.
//!
//! Local records are group-independent and are matched before filtering, so a
//! name that has a record here never reaches the blocklists.

const std = @import("std");
const Allocator = std.mem.Allocator;

const model = @import("../config/model.zig");
const header = @import("../dns/header.zig");
const name = @import("../dns/name.zig");
const packet = @import("../dns/packet.zig");
const question = @import("../dns/question.zig");
const types = @import("../dns/types.zig");
const address = @import("../platform/address.zig");

pub const Value = union(enum) { a: [4]u8, aaaa: [16]u8, cname: name.Name };

pub const Record = struct {
    /// Normalized owner name: lowercase, no trailing dot.
    owner: []const u8,
    value: Value,
    ttl: u32,
};

pub const Error = error{ OutOfMemory, BadRecordValue, BadRecordName, TooManyRecords };
pub const max_records: usize = 10_000;

pub const Records = struct {
    /// Sorted by (owner, rtype) so lookup is a binary search and the answer
    /// order for one name is stable across restarts.
    items: []const Record,
    /// One block holding every `Record.owner`; freed as a unit.
    owners: []const u8,

    pub const empty: Records = .{ .items = &.{}, .owners = &.{} };

    /// `gpa` owns the result; `deinit` frees it. Values are parsed here, once.
    /// A bad value is an error, not a skipped row — `validate.zig` already
    /// rejects these, so reaching one here means the database was edited behind
    /// nxdns's back and silence would make a record vanish with no signal.
    pub fn build(gpa: Allocator, rows: []const model.LocalRecord) Error!Records {
        if (rows.len == 0) return .empty;
        if (rows.len > max_records) return error.TooManyRecords;

        var owners: std.ArrayList(u8) = .empty;
        defer owners.deinit(gpa);
        var spans: std.ArrayList(Span) = .empty;
        defer spans.deinit(gpa);

        var buf: [types.max_name_len]u8 = undefined;
        for (rows) |row| {
            // A record whose owner can never be matched — a high byte, or the
            // root — is a configuration error, not a record that never answers.
            const owner = name.normalizeText(row.name, &buf) catch return error.BadRecordName;
            const value = try parseValue(row.rtype, row.value);
            try spans.append(gpa, .{
                .offset = owners.items.len,
                .len = owner.len,
                .value = value,
                .ttl = row.ttl,
            });
            try owners.appendSlice(gpa, owner);
        }

        const owner_bytes = try owners.toOwnedSlice(gpa);
        errdefer gpa.free(owner_bytes);

        const items = try gpa.alloc(Record, spans.items.len);
        for (items, spans.items) |*item, span| item.* = .{
            .owner = owner_bytes[span.offset..][0..span.len],
            .value = span.value,
            .ttl = span.ttl,
        };
        std.mem.sort(Record, items, {}, lessThan);

        return .{ .items = items, .owners = owner_bytes };
    }

    pub fn deinit(self: *Records, gpa: Allocator) void {
        gpa.free(self.items);
        gpa.free(self.owners);
        self.* = .empty;
    }

    /// All records for `domain` whose type matches `qtype`. `domain` must be
    /// normalized (lowercase, no trailing dot). An empty slice means the name
    /// has no local record of that type. Allocation-free.
    ///
    /// A CNAME answers every qtype and excludes every other type at the same
    /// name (RFC 1034 §3.6.2), so a name carrying one answers with the CNAME
    /// alone whatever else the row set holds.
    pub fn lookup(self: *const Records, domain: []const u8, qtype: types.Type) []const Record {
        const at_name = self.ownerRange(domain);
        if (at_name.len == 0) return at_name;

        const cnames = rankRun(at_name, rank_cname);
        if (cnames.len != 0) return cnames;

        return switch (qtype) {
            .a => rankRun(at_name, rank_a),
            .aaaa => rankRun(at_name, rank_aaaa),
            .any => at_name,
            else => at_name[0..0],
        };
    }

    /// True when the name has any local record of any type. The handler needs
    /// this to answer NODATA instead of forwarding a name nxdns owns.
    pub fn hasName(self: *const Records, domain: []const u8) bool {
        return self.ownerRange(domain).len != 0;
    }

    fn ownerRange(self: *const Records, domain: []const u8) []const Record {
        var low: usize = 0;
        var high: usize = self.items.len;
        while (low < high) {
            const mid = low + (high - low) / 2;
            if (std.mem.order(u8, self.items[mid].owner, domain) == .lt) {
                low = mid + 1;
            } else {
                high = mid;
            }
        }
        var end = low;
        while (end < self.items.len and std.mem.eql(u8, self.items[end].owner, domain)) end += 1;
        return self.items[low..end];
    }
};

/// Writes `records` as answers into a builder the caller has already
/// initialized with the request header and question. Mechanism only.
pub fn writeAnswers(
    b: *packet.ResponseBuilder,
    owner: name.Name,
    records: []const Record,
) packet.ResponseBuilder.Error!void {
    for (records) |rec| {
        switch (rec.value) {
            .a => |bytes| try b.addAnswer(owner, .a, .in, rec.ttl, &bytes),
            .aaaa => |bytes| try b.addAnswer(owner, .aaaa, .in, rec.ttl, &bytes),
            .cname => |target| try b.addAnswer(owner, .cname, .in, rec.ttl, target.wire()),
        }
    }
}

// ---------------------------------------------------------------------------
// Internals
// ---------------------------------------------------------------------------

/// `Record.owner` slices are cut only after the owner block stops growing, so
/// the build pass records offsets instead of pointers.
const Span = struct {
    offset: usize,
    len: usize,
    value: Value,
    ttl: u32,
};

const rank_a: u2 = 0;
const rank_aaaa: u2 = 1;
const rank_cname: u2 = 2;

fn rank(value: Value) u2 {
    return switch (value) {
        .a => rank_a,
        .aaaa => rank_aaaa,
        .cname => rank_cname,
    };
}

fn lessThan(_: void, a: Record, b: Record) bool {
    return switch (std.mem.order(u8, a.owner, b.owner)) {
        .lt => true,
        .gt => false,
        .eq => rank(a.value) < rank(b.value),
    };
}

/// The run of one rank inside a single name's records, which the (owner, rtype)
/// sort makes contiguous.
fn rankRun(records: []const Record, wanted: u2) []const Record {
    var start: usize = 0;
    while (start < records.len and rank(records[start].value) < wanted) start += 1;
    var end = start;
    while (end < records.len and rank(records[end].value) == wanted) end += 1;
    return records[start..end];
}

fn parseValue(rtype: model.RecordType, text: []const u8) error{BadRecordValue}!Value {
    switch (rtype) {
        .a => {
            const addr = address.NetAddress.parse(text) catch return error.BadRecordValue;
            return switch (addr) {
                .ip4 => |bytes| Value{ .a = bytes },
                .ip6 => return error.BadRecordValue,
            };
        },
        .aaaa => {
            const addr = address.NetAddress.parse(text) catch return error.BadRecordValue;
            return switch (addr) {
                .ip4 => return error.BadRecordValue,
                .ip6 => |bytes| Value{ .aaaa = bytes },
            };
        },
        .cname => {
            var buf: [types.max_name_len]u8 = undefined;
            const target = name.normalizeText(text, &buf) catch return error.BadRecordValue;
            return .{ .cname = name.fromText(target) catch return error.BadRecordValue };
        },
    }
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

const testing = std.testing;

const sample_rows = [_]model.LocalRecord{
    .{ .name = "nas.lan", .rtype = .a, .value = "192.168.1.5", .ttl = 120 },
    .{ .name = "nas.lan", .rtype = .aaaa, .value = "fd00::1", .ttl = 240 },
    .{ .name = "printer.lan", .rtype = .a, .value = "192.168.1.6", .ttl = 60 },
};

test "lookup returns the records of the queried type" {
    var records = try Records.build(testing.allocator, &sample_rows);
    defer records.deinit(testing.allocator);

    const a = records.lookup("nas.lan", .a);
    try testing.expectEqual(@as(usize, 1), a.len);
    try testing.expectEqualSlices(u8, &.{ 192, 168, 1, 5 }, &a[0].value.a);
    try testing.expectEqual(@as(u32, 120), a[0].ttl);

    const aaaa = records.lookup("nas.lan", .aaaa);
    try testing.expectEqual(@as(usize, 1), aaaa.len);
    try testing.expectEqual(@as(u32, 240), aaaa[0].ttl);

    try testing.expectEqual(@as(usize, 0), records.lookup("nas.lan", .mx).len);
    try testing.expectEqual(@as(usize, 0), records.lookup("other.lan", .a).len);
}

test "lookup returns the CNAME for every qtype" {
    const rows = [_]model.LocalRecord{
        .{ .name = "www.lan", .rtype = .cname, .value = "nas.lan", .ttl = 300 },
    };
    var records = try Records.build(testing.allocator, &rows);
    defer records.deinit(testing.allocator);

    for ([_]types.Type{ .a, .aaaa, .mx, .https, .any }) |qtype| {
        const found = records.lookup("www.lan", qtype);
        try testing.expectEqual(@as(usize, 1), found.len);
        try testing.expectEqualSlices(
            u8,
            (try name.fromText("nas.lan")).wire(),
            found[0].value.cname.wire(),
        );
    }
}

test "hasName covers every type at the name" {
    var records = try Records.build(testing.allocator, &sample_rows);
    defer records.deinit(testing.allocator);

    try testing.expect(records.hasName("nas.lan"));
    try testing.expect(records.hasName("printer.lan"));
    try testing.expect(!records.hasName("lan"));
    try testing.expect(!records.hasName("nas.lan.evil.net"));
}

test "two A records for one name come back in a stable order" {
    const rows = [_]model.LocalRecord{
        .{ .name = "nas.lan", .rtype = .a, .value = "192.168.1.5" },
        .{ .name = "nas.lan", .rtype = .a, .value = "192.168.1.6" },
    };

    var first = try Records.build(testing.allocator, &rows);
    defer first.deinit(testing.allocator);
    var second = try Records.build(testing.allocator, &rows);
    defer second.deinit(testing.allocator);

    const a = first.lookup("nas.lan", .a);
    const b = second.lookup("nas.lan", .a);
    try testing.expectEqual(@as(usize, 2), a.len);
    try testing.expectEqual(@as(usize, 2), b.len);
    try testing.expectEqualSlices(u8, &.{ 192, 168, 1, 5 }, &a[0].value.a);
    try testing.expectEqualSlices(u8, &.{ 192, 168, 1, 6 }, &a[1].value.a);
    for (a, b) |lhs, rhs| try testing.expectEqualSlices(u8, &lhs.value.a, &rhs.value.a);
}

test "uppercase and trailing-dot owners normalize to one key" {
    const rows = [_]model.LocalRecord{
        .{ .name = "NAS.Lan.", .rtype = .a, .value = "192.168.1.5" },
    };
    var records = try Records.build(testing.allocator, &rows);
    defer records.deinit(testing.allocator);

    try testing.expectEqualStrings("nas.lan", records.items[0].owner);
    try testing.expectEqual(@as(usize, 1), records.lookup("nas.lan", .a).len);
}

test "a bad A value is an error" {
    const rows = [_]model.LocalRecord{
        .{ .name = "nas.lan", .rtype = .a, .value = "::1" },
    };
    try testing.expectError(error.BadRecordValue, Records.build(testing.allocator, &rows));
}

test "a bad AAAA value is an error" {
    const rows = [_]model.LocalRecord{
        .{ .name = "nas.lan", .rtype = .aaaa, .value = "192.168.1.5" },
    };
    try testing.expectError(error.BadRecordValue, Records.build(testing.allocator, &rows));
}

test "a bad CNAME target is an error" {
    const rows = [_]model.LocalRecord{
        .{ .name = "www.lan", .rtype = .cname, .value = "nas..lan" },
    };
    try testing.expectError(error.BadRecordValue, Records.build(testing.allocator, &rows));
}

test "an unparseable owner is an error" {
    const rows = [_]model.LocalRecord{
        .{ .name = "nas..lan", .rtype = .a, .value = "192.168.1.5" },
    };
    try testing.expectError(error.BadRecordName, Records.build(testing.allocator, &rows));

    const root = [_]model.LocalRecord{
        .{ .name = ".", .rtype = .a, .value = "192.168.1.5" },
    };
    try testing.expectError(error.BadRecordName, Records.build(testing.allocator, &root));
}

test "too many rows is an error" {
    const rows = try testing.allocator.alloc(model.LocalRecord, max_records + 1);
    defer testing.allocator.free(rows);
    for (rows) |*row| row.* = .{ .name = "nas.lan", .rtype = .a, .value = "192.168.1.5" };
    try testing.expectError(error.TooManyRecords, Records.build(testing.allocator, rows));
}

test "an empty row set builds the empty table" {
    var records = try Records.build(testing.allocator, &[_]model.LocalRecord{});
    defer records.deinit(testing.allocator);

    try testing.expectEqual(@as(usize, 0), records.items.len);
    try testing.expect(!records.hasName("nas.lan"));
    try testing.expectEqual(@as(usize, 0), records.lookup("nas.lan", .a).len);
}

fn requestHeader() header.Header {
    return .{
        .id = 0x4242,
        .flags = .{
            .rcode = .no_error,
            .z = 0,
            .ra = false,
            .rd = true,
            .tc = false,
            .aa = false,
            .opcode = .query,
            .qr = false,
        },
        .qdcount = 1,
        .ancount = 0,
        .nscount = 0,
        .arcount = 0,
    };
}

test "writeAnswers emits records the parser reads back" {
    var records = try Records.build(testing.allocator, &sample_rows);
    defer records.deinit(testing.allocator);

    const owner = try name.fromText("nas.lan");
    const q: question.Question = .{ .name = owner, .qtype = .any, .qclass = .in };

    var buf: [512]u8 = undefined;
    var builder = try packet.ResponseBuilder.init(&buf, requestHeader(), q);
    try writeAnswers(&builder, owner, records.lookup("nas.lan", .any));
    const message = builder.finish();

    const parsed = try packet.parse(message);
    try testing.expectEqual(@as(u16, 2), parsed.header.ancount);

    var it = packet.answers(parsed);
    const first = (try it.next()).?;
    try testing.expectEqual(types.Type.a, first.rtype);
    try testing.expectEqual(@as(u16, @intFromEnum(types.Class.in)), first.class);
    try testing.expectEqual(@as(u32, 120), first.ttl);
    try testing.expectEqualSlices(u8, &.{ 192, 168, 1, 5 }, first.rdata.slice(parsed.bytes));

    const second = (try it.next()).?;
    try testing.expectEqual(types.Type.aaaa, second.rtype);
    try testing.expectEqual(@as(u32, 240), second.ttl);
    try testing.expectEqual(@as(usize, 16), second.rdata.len);

    try testing.expect((try it.next()) == null);
}

test "writeAnswers emits a CNAME in wire form" {
    const rows = [_]model.LocalRecord{
        .{ .name = "www.lan", .rtype = .cname, .value = "nas.lan", .ttl = 300 },
    };
    var records = try Records.build(testing.allocator, &rows);
    defer records.deinit(testing.allocator);

    const owner = try name.fromText("www.lan");
    const q: question.Question = .{ .name = owner, .qtype = .a, .qclass = .in };

    var buf: [512]u8 = undefined;
    var builder = try packet.ResponseBuilder.init(&buf, requestHeader(), q);
    try writeAnswers(&builder, owner, records.lookup("www.lan", .a));
    const message = builder.finish();

    const parsed = try packet.parse(message);
    try testing.expectEqual(@as(u16, 1), parsed.header.ancount);

    var it = packet.answers(parsed);
    const answer = (try it.next()).?;
    try testing.expectEqual(types.Type.cname, answer.rtype);
    try testing.expectEqual(@as(u32, 300), answer.ttl);
    try testing.expectEqualSlices(
        u8,
        (try name.fromText("nas.lan")).wire(),
        answer.rdata.slice(parsed.bytes),
    );
}

fn buildUnderFailure(gpa: Allocator) !void {
    var records = try Records.build(gpa, &sample_rows);
    defer records.deinit(gpa);
    try testing.expectEqual(@as(usize, 3), records.items.len);
}

test "build leaks nothing under allocation failure" {
    try testing.checkAllAllocationFailures(testing.allocator, buildUnderFailure, .{});
}