//! An immutable exact-match set of domain names, built from a compiled
//! blocklist body.
//!
//! The layout is a flat arena of length-prefixed names plus an open-addressed
//! table of `u32` offsets into it. No Bloom filter and no hash-only key set: a
//! false positive in a DNS sinkhole blocks a real domain for a real household
//! and is undebuggable from the outside, so every probe compares full bytes.
//!
//! Pure. Takes an allocator and bytes; no `std.Io`, no clock, no entropy
//! source — the hash seed arrives as a parameter.

const std = @import("std");

pub const DomainSet = struct {
    /// Length-prefixed lowercase names, back to back: [len: u8][bytes]…
    arena: []const u8,
    /// Open-addressed table of offsets into `arena`; `empty_slot` marks a hole.
    /// Length is always a power of two.
    index: []const u32,
    count: u32,
    seed: u64,

    pub const empty_slot: u32 = std.math.maxInt(u32);
    pub const max_count: u32 = 4_000_000;
    pub const max_arena_bytes: usize = 1 << 31;

    pub const Error = error{ OutOfMemory, TooManyDomains, SetTooLarge, NotSorted, NotLowercase };

    /// An empty set that owns nothing. `contains` on it is always false.
    pub const empty: DomainSet = .{
        .arena = &.{},
        .index = &.{},
        .count = 0,
        .seed = 0,
    };

    /// Builds from a compiled body: LF-separated, lowercase, sorted ascending,
    /// deduplicated, every line 1–255 bytes.
    ///
    /// The sortedness is verified, not assumed: a hand-edited or truncated file
    /// must fail loudly rather than produce a set that silently misses entries.
    /// Verification is one comparison per line. Because a valid line is at least
    /// one byte, comparing against an initial empty `prev` also rejects an empty
    /// line and a duplicate line — both are `error.NotSorted`, since neither is
    /// strictly ascending.
    ///
    /// A line longer than 255 bytes is `error.SetTooLarge`: the arena stores a
    /// `u8` length prefix and must not silently truncate, for the same reason
    /// the `u32` offsets bound `max_arena_bytes`.
    ///
    /// `seed` randomizes the hash. Query names are attacker-supplied, so a fixed
    /// seed would make probe-chain flooding computable offline.
    pub fn build(gpa: std.mem.Allocator, body: []const u8, seed: u64) Error!DomainSet {
        if (body.len == 0) return empty;
        if (body.len > max_arena_bytes) return error.SetTooLarge;

        var count: u32 = 0;
        var arena_len: usize = 0;
        {
            var prev: []const u8 = "";
            var it: LineIterator = .{ .rest = body };
            while (it.next()) |line| {
                for (line) |c| if (c >= 'A' and c <= 'Z') return error.NotLowercase;
                if (line.len > std.math.maxInt(u8)) return error.SetTooLarge;
                if (std.mem.order(u8, prev, line) != .lt) return error.NotSorted;
                if (count == max_count) return error.TooManyDomains;
                prev = line;
                count += 1;
                arena_len += 1 + line.len;
            }
        }

        const capacity = capacityFor(count);

        const arena = try gpa.alloc(u8, arena_len);
        errdefer gpa.free(arena);
        const index = try gpa.alloc(u32, capacity);
        @memset(index, empty_slot);

        const mask = capacity - 1;
        var write_at: usize = 0;
        var it: LineIterator = .{ .rest = body };
        while (it.next()) |line| {
            const offset: u32 = @intCast(write_at);
            arena[write_at] = @intCast(line.len);
            @memcpy(arena[write_at + 1 ..][0..line.len], line);
            write_at += 1 + line.len;

            var slot: usize = @as(usize, @truncate(std.hash.Wyhash.hash(seed, line))) & mask;
            while (index[slot] != empty_slot) slot = (slot + 1) & mask;
            index[slot] = offset;
        }

        return .{ .arena = arena, .index = index, .count = count, .seed = seed };
    }

    pub fn deinit(self: *DomainSet, gpa: std.mem.Allocator) void {
        gpa.free(self.arena);
        gpa.free(self.index);
        self.* = empty;
    }

    /// `domain` must be normalized (lowercase, no trailing dot). Allocation-free.
    pub fn contains(self: *const DomainSet, domain: []const u8) bool {
        if (self.index.len == 0) return false;
        const mask = self.index.len - 1;
        var slot: usize = @as(usize, @truncate(std.hash.Wyhash.hash(self.seed, domain))) & mask;
        while (true) {
            const offset = self.index[slot];
            if (offset == empty_slot) return false;
            const len = self.arena[offset];
            if (len == domain.len and std.mem.eql(u8, self.arena[offset + 1 ..][0..len], domain)) {
                return true;
            }
            slot = (slot + 1) & mask;
        }
    }

    /// Bytes held, for the memory report in `Snapshot.memoryBytes`.
    pub fn memoryBytes(self: *const DomainSet) usize {
        return self.arena.len + self.index.len * @sizeOf(u32);
    }
};

/// Smallest power of two at least `count * 4 / 3`, minimum 16. The load factor
/// stays at or under 0.75, which keeps at least one hole and therefore
/// terminates the probe loop in `contains`.
fn capacityFor(count: u32) usize {
    const wanted = (@as(u64, count) * 4 + 2) / 3;
    var capacity: usize = 16;
    while (capacity < wanted) capacity *= 2;
    return capacity;
}

/// Yields LF-separated lines. A final line without a trailing LF is yielded;
/// a trailing LF does not yield a final empty line.
const LineIterator = struct {
    rest: []const u8,

    fn next(self: *LineIterator) ?[]const u8 {
        if (self.rest.len == 0) return null;
        if (std.mem.indexOfScalar(u8, self.rest, '\n')) |nl| {
            defer self.rest = self.rest[nl + 1 ..];
            return self.rest[0..nl];
        }
        defer self.rest = self.rest[self.rest.len..];
        return self.rest;
    }
};

const testing = std.testing;

const small_body =
    "ads.example.com\n" ++
    "example.com\n" ++
    "tracker.example.org\n" ++
    "zzz.example.net\n";

test "build and contains over a small body" {
    var set = try DomainSet.build(testing.allocator, small_body, 0x1234);
    defer set.deinit(testing.allocator);

    try testing.expectEqual(@as(u32, 4), set.count);
    try testing.expect(set.contains("ads.example.com"));
    try testing.expect(set.contains("example.com"));
    try testing.expect(set.contains("tracker.example.org"));
    try testing.expect(set.contains("zzz.example.net"));

    // A prefix, a suffix, an uppercase spelling and the empty string.
    try testing.expect(!set.contains("example.co"));
    try testing.expect(!set.contains("example.com.evil.net"));
    try testing.expect(!set.contains("Example.com"));
    try testing.expect(!set.contains(""));
}

test "build rejects an unsorted body" {
    try testing.expectError(
        error.NotSorted,
        DomainSet.build(testing.allocator, "b.example.com\na.example.com\n", 0),
    );
}

test "build rejects an uppercase byte" {
    try testing.expectError(
        error.NotLowercase,
        DomainSet.build(testing.allocator, "a.example.com\nB.example.com\n", 0),
    );
}

test "build rejects a duplicate line" {
    try testing.expectError(
        error.NotSorted,
        DomainSet.build(testing.allocator, "a.example.com\na.example.com\n", 0),
    );
}

test "build rejects an over-long line" {
    var body: [300]u8 = undefined;
    @memset(&body, 'a');
    body[299] = '\n';
    try testing.expectError(error.SetTooLarge, DomainSet.build(testing.allocator, &body, 0));
}

test "contains does not depend on the seed" {
    var a = try DomainSet.build(testing.allocator, small_body, 0);
    defer a.deinit(testing.allocator);
    var b = try DomainSet.build(testing.allocator, small_body, 0xdead_beef_cafe_f00d);
    defer b.deinit(testing.allocator);

    var buf: [64]u8 = undefined;
    var i: usize = 0;
    while (i < 50) : (i += 1) {
        const probe = try std.fmt.bufPrint(&buf, "n{d}.example.com", .{i});
        try testing.expectEqual(a.contains(probe), b.contains(probe));
    }
    for ([_][]const u8{ "ads.example.com", "example.com", "zzz.example.net", "nope.test" }) |probe| {
        try testing.expectEqual(a.contains(probe), b.contains(probe));
    }
}

test "ten thousand names round-trip" {
    const gpa = testing.allocator;
    var body: std.ArrayList(u8) = .empty;
    defer body.deinit(gpa);

    // Fixed-width zero padding makes the generated order the sorted order.
    var line_buf: [64]u8 = undefined;
    var i: usize = 0;
    while (i < 10_000) : (i += 1) {
        try body.appendSlice(gpa, try std.fmt.bufPrint(&line_buf, "d{d:0>5}.example.com\n", .{i}));
    }

    var set = try DomainSet.build(gpa, body.items, 0x5eed);
    defer set.deinit(gpa);
    try testing.expectEqual(@as(u32, 10_000), set.count);

    var buf: [64]u8 = undefined;
    i = 0;
    while (i < 10_000) : (i += 1) {
        const nameStr = try std.fmt.bufPrint(&buf, "d{d:0>5}.example.com", .{i});
        try testing.expect(set.contains(nameStr));
    }
    try testing.expect(!set.contains("d10000.example.com"));
    try testing.expect(set.memoryBytes() < 2 * body.items.len);
}

fn buildUnderFailure(gpa: std.mem.Allocator) !void {
    var set = try DomainSet.build(gpa, small_body, 0x1234);
    defer set.deinit(gpa);
    try testing.expect(set.contains("example.com"));
}

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

test "the empty set owns nothing" {
    var set: DomainSet = .empty;
    try testing.expect(!set.contains("x"));
    try testing.expect(!set.contains(""));
    try testing.expectEqual(@as(usize, 0), set.memoryBytes());
    set.deinit(testing.allocator);
    try testing.expect(!set.contains("x"));
}

test "an empty body builds the empty set" {
    var set = try DomainSet.build(testing.allocator, "", 7);
    defer set.deinit(testing.allocator);
    try testing.expectEqual(@as(u32, 0), set.count);
    try testing.expect(!set.contains("x"));
}

test "a final line without a trailing newline is kept" {
    var set = try DomainSet.build(testing.allocator, "a.example.com\nb.example.com", 0);
    defer set.deinit(testing.allocator);
    try testing.expectEqual(@as(u32, 2), set.count);
    try testing.expect(set.contains("b.example.com"));
}

test "capacityFor keeps the load factor at or under three quarters" {
    try testing.expectEqual(@as(usize, 16), capacityFor(0));
    try testing.expectEqual(@as(usize, 16), capacityFor(12));
    try testing.expectEqual(@as(usize, 32), capacityFor(13));
    try testing.expectEqual(@as(usize, 2 << 20), capacityFor(1_000_000));
}
