//! IP address and prefix values for nxdns. Pure: no `std.Io` operations, only
//! conversions to and from `std.Io.net.IpAddress`.

const std = @import("std");
const net = std.Io.net;
const Writer = std.Io.Writer;

pub const NetAddress = union(enum) {
    ip4: [4]u8,
    ip6: [16]u8,

    /// Tag byte (4 or 6) followed by the address bytes, zero-padded for ip4.
    pub const Key = [17]u8;

    pub const ParseError = error{InvalidAddress};

    /// Accepts "1.2.3.4" and "fd00::1". No port, no brackets, no scope id.
    /// An IPv4-mapped IPv6 literal normalizes to `.ip4`.
    pub fn parse(text: []const u8) ParseError!NetAddress {
        const addr = net.IpAddress.parse(text, 0) catch return error.InvalidAddress;
        return fromIp(addr);
    }

    /// IPv4 in dotted-decimal, IPv6 per RFC 5952.
    pub fn format(self: NetAddress, w: *Writer) Writer.Error!void {
        switch (self) {
            .ip4 => |b| try w.print("{d}.{d}.{d}.{d}", .{ b[0], b[1], b[2], b[3] }),
            .ip6 => |b| try formatIp6(b, w),
        }
    }

    pub fn key(self: NetAddress) Key {
        var out: Key = @splat(0);
        switch (self) {
            .ip4 => |b| {
                out[0] = 4;
                @memcpy(out[1..5], &b);
            },
            .ip6 => |b| {
                out[0] = 6;
                @memcpy(out[1..17], &b);
            },
        }
        return out;
    }

    /// Drops the port. An IPv4-mapped IPv6 address (::ffff:a.b.c.d) yields `.ip4`.
    pub fn fromIp(addr: net.IpAddress) NetAddress {
        return switch (addr) {
            .ip4 => |a| .{ .ip4 = a.bytes },
            .ip6 => |a| if (net.Ip4Address.fromIp6(a)) |mapped|
                .{ .ip4 = mapped.bytes }
            else
                .{ .ip6 = a.bytes },
        };
    }

    pub fn toIp(self: NetAddress, port: u16) net.IpAddress {
        return switch (self) {
            .ip4 => |b| .{ .ip4 = .{ .bytes = b, .port = port } },
            .ip6 => |b| .{ .ip6 = .{ .bytes = b, .port = port } },
        };
    }

    pub fn eql(a: NetAddress, b: NetAddress) bool {
        return switch (a) {
            .ip4 => |x| switch (b) {
                .ip4 => |y| std.mem.eql(u8, &x, &y),
                .ip6 => false,
            },
            .ip6 => |x| switch (b) {
                .ip4 => false,
                .ip6 => |y| std.mem.eql(u8, &x, &y),
            },
        };
    }

    fn bytes(self: *const NetAddress) []const u8 {
        return switch (self.*) {
            .ip4 => |*b| b,
            .ip6 => |*b| b,
        };
    }

    fn familyBits(self: NetAddress) u8 {
        return switch (self) {
            .ip4 => 32,
            .ip6 => 128,
        };
    }
};

fn formatIp6(addr: [16]u8, w: *Writer) Writer.Error!void {
    var parts: [8]u16 = undefined;
    for (&parts, 0..) |*part, i| {
        part.* = std.mem.readInt(u16, addr[i * 2 ..][0..2], .big);
    }

    // RFC 5952 4.2: compress the longest run of zero groups, earliest run on a
    // tie, and never a run of only one group.
    var run_start: usize = 0;
    var run_len: usize = 0;
    var current_start: usize = 0;
    var current_len: usize = 0;
    for (parts, 0..) |part, i| {
        if (part != 0) {
            current_len = 0;
            continue;
        }
        if (current_len == 0) current_start = i;
        current_len += 1;
        if (current_len > run_len) {
            run_start = current_start;
            run_len = current_len;
        }
    }
    if (run_len < 2) run_len = 0;

    var i: usize = 0;
    var at_start = true;
    while (i < parts.len) {
        if (run_len != 0 and i == run_start) {
            try w.writeAll("::");
            i += run_len;
            at_start = true;
            continue;
        }
        if (!at_start) try w.writeByte(':');
        try w.print("{x}", .{parts[i]});
        at_start = false;
        i += 1;
    }
}

pub const Prefix = struct {
    /// Host bits are zeroed by `parse`.
    addr: NetAddress,
    bits: u8,

    pub const ParseError = error{InvalidPrefix};

    /// Accepts "192.168.1.0/24" and "fd00:abcd::/48".
    pub fn parse(text: []const u8) ParseError!Prefix {
        const slash = std.mem.findScalar(u8, text, '/') orelse return error.InvalidPrefix;
        const addr = NetAddress.parse(text[0..slash]) catch return error.InvalidPrefix;
        const bits = std.fmt.parseInt(u8, text[slash + 1 ..], 10) catch return error.InvalidPrefix;
        if (bits > addr.familyBits()) return error.InvalidPrefix;

        var masked = addr;
        switch (masked) {
            .ip4 => |*b| maskBytes(b, bits),
            .ip6 => |*b| maskBytes(b, bits),
        }
        return .{ .addr = masked, .bits = bits };
    }

    /// A family mismatch is never a match.
    pub fn contains(self: Prefix, addr: NetAddress) bool {
        if (std.meta.activeTag(self.addr) != std.meta.activeTag(addr)) return false;
        const net_bytes = self.addr.bytes();
        const addr_bytes = addr.bytes();
        var remaining = self.bits;
        var i: usize = 0;
        while (remaining >= 8) : (i += 1) {
            if (net_bytes[i] != addr_bytes[i]) return false;
            remaining -= 8;
        }
        if (remaining == 0) return true;
        const mask = ~(@as(u8, 0xff) >> @intCast(remaining));
        return (net_bytes[i] & mask) == (addr_bytes[i] & mask);
    }

    pub fn format(self: Prefix, w: *Writer) Writer.Error!void {
        try self.addr.format(w);
        try w.print("/{d}", .{self.bits});
    }
};

fn maskBytes(addr: []u8, bits: u8) void {
    for (addr, 0..) |*byte, i| {
        const offset = i * 8;
        if (bits >= offset + 8) continue;
        if (bits <= offset) {
            byte.* = 0;
            continue;
        }
        byte.* &= ~(@as(u8, 0xff) >> @intCast(bits - offset));
    }
}

/// Longest-prefix winner; ties broken by lower `priority` value. Returns null
/// when nothing matches. `T` must have fields `prefix: Prefix` and
/// `priority: i64`.
pub fn matchLongest(comptime T: type, entries: []const T, addr: NetAddress) ?*const T {
    var best: ?*const T = null;
    for (entries) |*entry| {
        if (!entry.prefix.contains(addr)) continue;
        const current = best orelse {
            best = entry;
            continue;
        };
        if (entry.prefix.bits > current.prefix.bits) {
            best = entry;
        } else if (entry.prefix.bits == current.prefix.bits and entry.priority < current.priority) {
            best = entry;
        }
    }
    return best;
}

const testing = std.testing;

fn expectFormat(expected: []const u8, value: anytype) !void {
    var buf: [64]u8 = undefined;
    var w = Writer.fixed(&buf);
    try value.format(&w);
    try testing.expectEqualStrings(expected, w.buffered());
}

fn expectRoundTrip(text: []const u8) !void {
    try expectFormat(text, try NetAddress.parse(text));
}

test "round trip of canonical addresses" {
    try expectRoundTrip("192.168.1.1");
    try expectRoundTrip("0.0.0.0");
    try expectRoundTrip("255.255.255.255");
    try expectRoundTrip("fd00::1");
    try expectRoundTrip("::");
    try expectRoundTrip("::1");
    try expectRoundTrip("2001:db8::8:800:200c:417a");
    try expectRoundTrip("2001:db8::");
}

test "rfc 5952 formatting rules" {
    // Longest run wins; earliest run wins a tie; uppercase input lowercases.
    try expectFormat("2001:db8::1:0:0:1", try NetAddress.parse("2001:0DB8:0:0:1::1"));
    // A single zero group is never compressed.
    try expectFormat("2001:db8:0:1:1:1:1:1", try NetAddress.parse("2001:db8:0:1:1:1:1:1"));
    // Leading zeros in a group are suppressed.
    try expectFormat("2001:db8:aaaa:bbbb:cccc:dddd:eeee:1", try NetAddress.parse("2001:0db8:aaaa:bbbb:cccc:dddd:eeee:0001"));
    // The longest run is compressed even when a shorter run comes first.
    try expectFormat("2001:0:0:1::1", try NetAddress.parse("2001:0:0:1:0:0:0:1"));
    // A trailing run compresses to a trailing "::".
    try expectFormat("2001:db8::", try NetAddress.parse("2001:0db8:0000:0000:0000:0000:0000:0000"));
    // Hex digits are lowercase.
    try expectFormat("fe80::abcd:ef01", try NetAddress.parse("FE80::ABCD:EF01"));
}

test "parse rejects malformed input" {
    try testing.expectError(error.InvalidAddress, NetAddress.parse(""));
    try testing.expectError(error.InvalidAddress, NetAddress.parse("1.2.3"));
    try testing.expectError(error.InvalidAddress, NetAddress.parse("1.2.3.4.5"));
    try testing.expectError(error.InvalidAddress, NetAddress.parse("256.1.1.1"));
    try testing.expectError(error.InvalidAddress, NetAddress.parse("1.2.3.4:53"));
    try testing.expectError(error.InvalidAddress, NetAddress.parse("[fd00::1]"));
    try testing.expectError(error.InvalidAddress, NetAddress.parse("[fd00::1]:853"));
    try testing.expectError(error.InvalidAddress, NetAddress.parse("fd00::1%eth0"));
    try testing.expectError(error.InvalidAddress, NetAddress.parse("fd00:::1"));
    try testing.expectError(error.InvalidAddress, NetAddress.parse("nonsense"));
}

test "key encodes family and bytes" {
    const v4 = try NetAddress.parse("192.168.1.1");
    var expected_v4: NetAddress.Key = @splat(0);
    expected_v4[0] = 4;
    expected_v4[1] = 192;
    expected_v4[2] = 168;
    expected_v4[3] = 1;
    expected_v4[4] = 1;
    try testing.expectEqualSlices(u8, &expected_v4, &v4.key());

    const v6 = try NetAddress.parse("fd00::1");
    const k6 = v6.key();
    try testing.expectEqual(@as(u8, 6), k6[0]);
    try testing.expectEqual(@as(u8, 0xfd), k6[1]);
    try testing.expectEqual(@as(u8, 1), k6[16]);

    // Distinct families never collide.
    try testing.expect(!std.mem.eql(u8, &v4.key(), &v6.key()));
}

test "fromIp normalizes ipv4-mapped ipv6" {
    const mapped = try net.IpAddress.parse("::ffff:192.168.1.1", 853);
    try testing.expectEqual(net.IpAddress.Family.ip6, std.meta.activeTag(mapped));

    const addr = NetAddress.fromIp(mapped);
    try testing.expectEqual(NetAddress.ip4, std.meta.activeTag(addr));
    try testing.expect(addr.eql(.{ .ip4 = .{ 192, 168, 1, 1 } }));
    try expectFormat("192.168.1.1", addr);

    // "::" is not a mapped address.
    const any6 = try net.IpAddress.parse("::", 0);
    try testing.expectEqual(NetAddress.ip6, std.meta.activeTag(NetAddress.fromIp(any6)));
}

test "toIp restores family and sets port" {
    const v4 = (try NetAddress.parse("10.0.0.1")).toIp(53);
    try testing.expectEqual(@as(u16, 53), v4.getPort());
    try testing.expectEqualSlices(u8, &.{ 10, 0, 0, 1 }, &v4.ip4.bytes);

    const v6 = (try NetAddress.parse("fd00::1")).toIp(853);
    try testing.expectEqual(@as(u16, 853), v6.getPort());
    try testing.expect(NetAddress.fromIp(v6).eql(try NetAddress.parse("fd00::1")));
}

test "eql compares family and bytes" {
    const a = try NetAddress.parse("192.168.1.1");
    const b = try NetAddress.parse("192.168.1.2");
    try testing.expect(a.eql(a));
    try testing.expect(!a.eql(b));
    try testing.expect(!a.eql(try NetAddress.parse("::")));
    try testing.expect((try NetAddress.parse("fd00::1")).eql(try NetAddress.parse("fd00:0:0:0:0:0:0:1")));
}

test "prefix parse zeroes host bits" {
    try expectFormat("192.168.1.0/24", try Prefix.parse("192.168.1.55/24"));
    try expectFormat("0.0.0.0/0", try Prefix.parse("192.168.1.55/0"));
    try expectFormat("10.128.0.0/9", try Prefix.parse("10.255.3.4/9"));
    try expectFormat("192.168.1.55/32", try Prefix.parse("192.168.1.55/32"));
    try expectFormat("fd00:abcd::/48", try Prefix.parse("fd00:abcd:0:1234::5/48"));
    try expectFormat("::/0", try Prefix.parse("2001:db8::1/0"));
}

test "prefix parse rejects malformed input" {
    try testing.expectError(error.InvalidPrefix, Prefix.parse("192.168.1.0"));
    try testing.expectError(error.InvalidPrefix, Prefix.parse("192.168.1.0/"));
    try testing.expectError(error.InvalidPrefix, Prefix.parse("192.168.1.0/33"));
    try testing.expectError(error.InvalidPrefix, Prefix.parse("fd00::/129"));
    try testing.expectError(error.InvalidPrefix, Prefix.parse("/24"));
    try testing.expectError(error.InvalidPrefix, Prefix.parse("192.168.1.0/x"));
}

test "prefix contains" {
    const p24 = try Prefix.parse("192.168.1.0/24");
    try testing.expect(p24.contains(try NetAddress.parse("192.168.1.5")));
    try testing.expect(p24.contains(try NetAddress.parse("192.168.1.0")));
    try testing.expect(p24.contains(try NetAddress.parse("192.168.1.255")));
    try testing.expect(!p24.contains(try NetAddress.parse("192.168.2.5")));

    // Family mismatch is false in both directions.
    try testing.expect(!p24.contains(try NetAddress.parse("fd00::1")));
    const p6 = try Prefix.parse("fd00::/8");
    try testing.expect(!p6.contains(try NetAddress.parse("192.168.1.5")));
    try testing.expect(p6.contains(try NetAddress.parse("fdff::abcd")));
    try testing.expect(!p6.contains(try NetAddress.parse("fe80::1")));

    // /0 contains everything of its family, nothing of the other.
    const any4 = try Prefix.parse("0.0.0.0/0");
    try testing.expect(any4.contains(try NetAddress.parse("8.8.8.8")));
    try testing.expect(any4.contains(try NetAddress.parse("0.0.0.0")));
    try testing.expect(!any4.contains(try NetAddress.parse("::")));
    const any6 = try Prefix.parse("::/0");
    try testing.expect(any6.contains(try NetAddress.parse("2001:db8::1")));
    try testing.expect(!any6.contains(try NetAddress.parse("8.8.8.8")));

    // Single-host prefixes.
    const host4 = try Prefix.parse("10.1.2.3/32");
    try testing.expect(host4.contains(try NetAddress.parse("10.1.2.3")));
    try testing.expect(!host4.contains(try NetAddress.parse("10.1.2.4")));
    const host6 = try Prefix.parse("fd00::1/128");
    try testing.expect(host6.contains(try NetAddress.parse("fd00::1")));
    try testing.expect(!host6.contains(try NetAddress.parse("fd00::2")));

    // Non-byte-aligned boundary.
    const p9 = try Prefix.parse("10.128.0.0/9");
    try testing.expect(p9.contains(try NetAddress.parse("10.255.255.255")));
    try testing.expect(!p9.contains(try NetAddress.parse("10.127.255.255")));
}

const Rule = struct {
    prefix: Prefix,
    priority: i64,
    name: []const u8,
};

test "matchLongest prefers longer prefixes" {
    const rules = [_]Rule{
        .{ .prefix = try Prefix.parse("192.168.0.0/16"), .priority = 0, .name = "wide" },
        .{ .prefix = try Prefix.parse("192.168.1.0/24"), .priority = 0, .name = "narrow" },
        .{ .prefix = try Prefix.parse("0.0.0.0/0"), .priority = 0, .name = "default" },
    };
    const hit = matchLongest(Rule, &rules, try NetAddress.parse("192.168.1.5")).?;
    try testing.expectEqualStrings("narrow", hit.name);

    const wide = matchLongest(Rule, &rules, try NetAddress.parse("192.168.2.5")).?;
    try testing.expectEqualStrings("wide", wide.name);

    const fallback = matchLongest(Rule, &rules, try NetAddress.parse("8.8.8.8")).?;
    try testing.expectEqualStrings("default", fallback.name);
}

test "matchLongest breaks ties by lower priority" {
    const rules = [_]Rule{
        .{ .prefix = try Prefix.parse("192.168.1.0/24"), .priority = 7, .name = "later" },
        .{ .prefix = try Prefix.parse("192.168.1.0/24"), .priority = -3, .name = "winner" },
        .{ .prefix = try Prefix.parse("192.168.1.0/24"), .priority = 2, .name = "middle" },
    };
    const hit = matchLongest(Rule, &rules, try NetAddress.parse("192.168.1.5")).?;
    try testing.expectEqualStrings("winner", hit.name);
}

test "matchLongest returns null when nothing matches" {
    const rules = [_]Rule{
        .{ .prefix = try Prefix.parse("192.168.1.0/24"), .priority = 0, .name = "a" },
        .{ .prefix = try Prefix.parse("fd00::/8"), .priority = 0, .name = "b" },
    };
    try testing.expect(matchLongest(Rule, &rules, try NetAddress.parse("10.0.0.1")) == null);
    try testing.expect(matchLongest(Rule, &rules, try NetAddress.parse("2001:db8::1")) == null);
    try testing.expect(matchLongest(Rule, &[_]Rule{}, try NetAddress.parse("10.0.0.1")) == null);
}
