// ©AngelaMos | 2026
// targets.zig

const std = @import("std");
const numtheory = @import("numtheory");
const netutil = @import("netutil");

pub const Range = struct {
    start: u32,
    end: u32,

    pub fn count(self: Range) u64 {
        return @as(u64, self.end - self.start) + 1;
    }
};

const reserved = [_]Range{
    .{ .start = 0x00000000, .end = 0x00ffffff },
    .{ .start = 0x0a000000, .end = 0x0affffff },
    .{ .start = 0x64400000, .end = 0x647fffff },
    .{ .start = 0x7f000000, .end = 0x7fffffff },
    .{ .start = 0xa9fe0000, .end = 0xa9feffff },
    .{ .start = 0xac100000, .end = 0xac1fffff },
    .{ .start = 0xc0000000, .end = 0xc00000ff },
    .{ .start = 0xc0000200, .end = 0xc00002ff },
    .{ .start = 0xc0586300, .end = 0xc05863ff },
    .{ .start = 0xc0a80000, .end = 0xc0a8ffff },
    .{ .start = 0xc6120000, .end = 0xc613ffff },
    .{ .start = 0xc6336400, .end = 0xc63364ff },
    .{ .start = 0xcb007100, .end = 0xcb0071ff },
    .{ .start = 0xe0000000, .end = 0xefffffff },
    .{ .start = 0xf0000000, .end = 0xffffffff },
};

pub fn parseCidr(text: []const u8) !Range {
    const slash = std.mem.indexOfScalar(u8, text, '/') orelse return error.InvalidCidr;
    const addr_text = text[0..slash];
    const prefix = std.fmt.parseInt(u6, text[slash + 1 ..], 10) catch return error.InvalidCidr;
    if (prefix > 32) return error.InvalidCidr;

    var base: u32 = 0;
    var octets: usize = 0;
    var it = std.mem.splitScalar(u8, addr_text, '.');
    while (it.next()) |part| {
        if (octets == 4) return error.InvalidCidr;
        const octet = std.fmt.parseInt(u8, part, 10) catch return error.InvalidCidr;
        base = (base << 8) | octet;
        octets += 1;
    }
    if (octets != 4) return error.InvalidCidr;

    const host_bits: u6 = @intCast(32 - @as(u32, prefix));
    if (host_bits == 32) return .{ .start = 0, .end = 0xffffffff };
    const sh: u5 = @intCast(host_bits);
    const span: u32 = (@as(u32, 1) << sh) - 1;
    const start = base & ~span;
    return .{ .start = start, .end = start | span };
}

pub fn isReserved(ip: u32) bool {
    var lo: usize = 0;
    var hi: usize = reserved.len;
    while (lo < hi) {
        const mid = lo + (hi - lo) / 2;
        if (ip < reserved[mid].start) {
            hi = mid;
        } else if (ip > reserved[mid].end) {
            lo = mid + 1;
        } else return true;
    }
    return false;
}

fn subtractReserved(allocator: std.mem.Allocator, acc: *std.ArrayList(Range), r: Range) !void {
    var pending: std.ArrayList(Range) = .empty;
    defer pending.deinit(allocator);
    try pending.append(allocator, r);
    for (reserved) |res| {
        var next: std.ArrayList(Range) = .empty;
        errdefer next.deinit(allocator);
        for (pending.items) |cur| {
            if (res.end < cur.start or res.start > cur.end) {
                try next.append(allocator, cur);
                continue;
            }
            if (cur.start < res.start) try next.append(allocator, .{ .start = cur.start, .end = res.start - 1 });
            if (cur.end > res.end) try next.append(allocator, .{ .start = res.end + 1, .end = cur.end });
        }
        pending.deinit(allocator);
        pending = next;
    }
    for (pending.items) |s| try acc.append(allocator, s);
}

pub const IpPicker = struct {
    allocator: std.mem.Allocator,
    ranges: []Range,
    prefix: []u64,
    count: u64,

    pub fn build(allocator: std.mem.Allocator, user: []const Range) !IpPicker {
        var acc: std.ArrayList(Range) = .empty;
        defer acc.deinit(allocator);
        for (user) |r| try subtractReserved(allocator, &acc, r);
        std.mem.sort(Range, acc.items, {}, struct {
            fn lt(_: void, a: Range, b: Range) bool {
                return a.start < b.start;
            }
        }.lt);

        const ranges = try allocator.dupe(Range, acc.items);
        errdefer allocator.free(ranges);
        const prefix = try allocator.alloc(u64, ranges.len + 1);
        var total: u64 = 0;
        for (ranges, 0..) |r, k| {
            prefix[k] = total;
            total += r.count();
        }
        prefix[ranges.len] = total;
        return .{ .allocator = allocator, .ranges = ranges, .prefix = prefix, .count = total };
    }

    pub fn deinit(self: *IpPicker) void {
        self.allocator.free(self.ranges);
        self.allocator.free(self.prefix);
    }

    pub fn at(self: IpPicker, index: u64) u32 {
        std.debug.assert(index < self.count);
        var lo: usize = 0;
        var hi: usize = self.ranges.len;
        while (lo + 1 < hi) {
            const mid = lo + (hi - lo) / 2;
            if (self.prefix[mid] <= index) lo = mid else hi = mid;
        }
        const offset: u32 = @intCast(index - self.prefix[lo]);
        return self.ranges[lo].start + offset;
    }
};

pub const Target = struct {
    ip: u32,
    port: u16,
};

pub const Engine = struct {
    picker: IpPicker,
    ports: []u16,
    num_ports: u64,
    total: u64,
    prime: u64,
    generator: u64,
    current: u64,
    steps_left: u64,

    pub fn init(allocator: std.mem.Allocator, cidrs: []const Range, ports: []const u16, seed: u64) !Engine {
        return initShard(allocator, cidrs, ports, seed, 1, 0);
    }

    pub fn initShard(
        allocator: std.mem.Allocator,
        cidrs: []const Range,
        ports: []const u16,
        seed: u64,
        num_shards: u64,
        shard_id: u64,
    ) !Engine {
        var picker = try IpPicker.build(allocator, cidrs);
        errdefer picker.deinit();
        const ports_copy = try allocator.dupe(u16, ports);
        errdefer allocator.free(ports_copy);

        const num_ports: u64 = @intCast(ports.len);
        const total = picker.count * num_ports;
        const prime = numtheory.smallestPrimeAbove(total);
        const order = prime - 1;
        if (num_shards == 0 or shard_id >= num_shards or num_shards > order) return error.InvalidShardCount;

        var prng = std.Random.DefaultPrng.init(seed);
        const rand = prng.random();
        const generator = numtheory.findPrimitiveRoot(prime, rand);
        const start = rand.intRangeAtMost(u64, 1, prime - 1);

        const chunk = order / num_shards;
        const begin = shard_id * chunk;
        const my_steps = if (shard_id == num_shards - 1) order - begin else chunk;
        const offset = numtheory.modExp(generator, begin, prime);
        const current = numtheory.mulMod(start, offset, prime);

        return .{
            .picker = picker,
            .ports = ports_copy,
            .num_ports = num_ports,
            .total = total,
            .prime = prime,
            .generator = generator,
            .current = current,
            .steps_left = my_steps,
        };
    }

    pub fn deinit(self: *Engine) void {
        const allocator = self.picker.allocator;
        self.picker.deinit();
        allocator.free(self.ports);
    }

    pub fn next(self: *Engine) ?Target {
        while (self.steps_left > 0) {
            self.current = numtheory.mulMod(self.current, self.generator, self.prime);
            self.steps_left -= 1;
            const idx = self.current;
            if (idx >= 1 and idx <= self.total) {
                const idx0 = idx - 1;
                const ip_pos = idx0 / self.num_ports;
                const port_pos = idx0 % self.num_ports;
                return .{ .ip = self.picker.at(ip_pos), .port = self.ports[@intCast(port_pos)] };
            }
        }
        return null;
    }
};

pub const default_max_hosts6: u64 = 1 << 20;

pub const Cidr6 = struct {
    base: [16]u8,
    prefix: u8,
};

const Reserved6 = struct { prefix: [16]u8, bits: u8 };

const reserved6 = [_]Reserved6{
    .{ .prefix = [_]u8{0} ** 16, .bits = 128 },
    .{ .prefix = [_]u8{0} ** 15 ++ [_]u8{1}, .bits = 128 },
    .{ .prefix = [_]u8{0} ** 10 ++ [_]u8{ 0xff, 0xff } ++ [_]u8{0} ** 4, .bits = 96 },
    .{ .prefix = [_]u8{ 0x01, 0x00 } ++ [_]u8{0} ** 14, .bits = 8 },
    .{ .prefix = [_]u8{ 0x20, 0x01, 0x0d, 0xb8 } ++ [_]u8{0} ** 12, .bits = 32 },
    .{ .prefix = [_]u8{ 0xfc, 0x00 } ++ [_]u8{0} ** 14, .bits = 7 },
    .{ .prefix = [_]u8{ 0xfe, 0x80 } ++ [_]u8{0} ** 14, .bits = 10 },
    .{ .prefix = [_]u8{ 0xff, 0x00 } ++ [_]u8{0} ** 14, .bits = 8 },
};

fn inPrefix6(addr: [16]u8, prefix: [16]u8, bits: u8) bool {
    const full = bits / 8;
    for (0..full) |i| if (addr[i] != prefix[i]) return false;
    const rem: u3 = @intCast(bits % 8);
    if (rem != 0) {
        const mask: u8 = @as(u8, 0xff) << @intCast(8 - @as(u4, rem));
        if ((addr[full] & mask) != (prefix[full] & mask)) return false;
    }
    return true;
}

pub fn isReserved6(addr: [16]u8) bool {
    for (reserved6) |res| {
        if (inPrefix6(addr, res.prefix, res.bits)) return true;
    }
    return false;
}

fn maskAddr6(addr: [16]u8, prefix: u8) [16]u8 {
    var out = addr;
    var bit: usize = prefix;
    while (bit < 128) : (bit += 1) {
        const byte = bit / 8;
        const off: u3 = @intCast(7 - (bit % 8));
        out[byte] &= ~(@as(u8, 1) << off);
    }
    return out;
}

pub fn parseCidr6(text: []const u8) !Cidr6 {
    const slash = std.mem.indexOfScalar(u8, text, '/') orelse return error.InvalidCidr;
    const addr = try netutil.parseIpv6(text[0..slash]);
    const prefix = std.fmt.parseInt(u8, text[slash + 1 ..], 10) catch return error.InvalidCidr;
    if (prefix > 128) return error.InvalidCidr;
    if (prefix == 0) return error.PrefixTooLarge;
    return .{ .base = maskAddr6(addr, prefix), .prefix = prefix };
}

pub const Target6 = struct {
    addr: [16]u8,
    port: u16,
};

pub const Engine6 = struct {
    base: [16]u8,
    ports: []u16,
    num_ports: u64,
    host_count: u64,
    total: u64,
    prime: u64,
    generator: u64,
    current: u64,
    steps_left: u64,
    allocator: std.mem.Allocator,

    pub fn init(allocator: std.mem.Allocator, cidr: Cidr6, ports: []const u16, seed: u64, max_hosts: u64) !Engine6 {
        const host_bits: u8 = 128 - cidr.prefix;
        if (host_bits >= 64) return error.PrefixTooLarge;
        const host_count: u64 = if (host_bits == 0) 1 else (@as(u64, 1) << @intCast(host_bits));
        if (host_count > max_hosts) return error.PrefixTooLarge;

        const ports_copy = try allocator.dupe(u16, ports);
        errdefer allocator.free(ports_copy);
        const num_ports: u64 = @intCast(ports.len);
        if (num_ports == 0 or host_count > std.math.maxInt(u64) / num_ports) return error.PrefixTooLarge;
        const total = host_count * num_ports;
        const prime = numtheory.smallestPrimeAbove(total);

        var prng = std.Random.DefaultPrng.init(seed);
        const rand = prng.random();
        const generator = numtheory.findPrimitiveRoot(prime, rand);
        const start = rand.intRangeAtMost(u64, 1, prime - 1);

        return .{
            .base = cidr.base,
            .ports = ports_copy,
            .num_ports = num_ports,
            .host_count = host_count,
            .total = total,
            .prime = prime,
            .generator = generator,
            .current = start,
            .steps_left = prime - 1,
            .allocator = allocator,
        };
    }

    pub fn deinit(self: *Engine6) void {
        self.allocator.free(self.ports);
    }

    fn addrAt(self: *const Engine6, host_index: u64) [16]u8 {
        var addr = self.base;
        const lo = std.mem.readInt(u64, addr[8..16], .big);
        std.mem.writeInt(u64, addr[8..16], lo | host_index, .big);
        return addr;
    }

    pub fn next(self: *Engine6) ?Target6 {
        while (self.steps_left > 0) {
            self.current = numtheory.mulMod(self.current, self.generator, self.prime);
            self.steps_left -= 1;
            const idx = self.current;
            if (idx >= 1 and idx <= self.total) {
                const idx0 = idx - 1;
                const host_pos = idx0 / self.num_ports;
                const port_pos = idx0 % self.num_ports;
                const addr = self.addrAt(host_pos);
                if (isReserved6(addr)) continue;
                return .{ .addr = addr, .port = self.ports[@intCast(port_pos)] };
            }
        }
        return null;
    }
};

test "parseCidr yields the right range and count" {
    const a = try parseCidr("10.0.0.0/24");
    try std.testing.expectEqual(@as(u32, 0x0a000000), a.start);
    try std.testing.expectEqual(@as(u32, 0x0a0000ff), a.end);
    try std.testing.expectEqual(@as(u64, 256), a.count());

    const b = try parseCidr("192.168.1.0/30");
    try std.testing.expectEqual(@as(u64, 4), b.count());

    const h = try parseCidr("8.8.8.8/32");
    try std.testing.expectEqual(@as(u32, 0x08080808), h.start);
    try std.testing.expectEqual(@as(u64, 1), h.count());

    try std.testing.expectError(error.InvalidCidr, parseCidr("999.0.0.0/8"));
    try std.testing.expectError(error.InvalidCidr, parseCidr("10.0.0.0/33"));
}

test "isReserved flags RFC 6890 space, passes public IPs" {
    try std.testing.expect(isReserved((try parseCidr("127.0.0.1/32")).start));
    try std.testing.expect(isReserved((try parseCidr("10.1.2.3/32")).start));
    try std.testing.expect(isReserved((try parseCidr("192.168.1.1/32")).start));
    try std.testing.expect(isReserved((try parseCidr("169.254.5.5/32")).start));
    try std.testing.expect(isReserved((try parseCidr("224.0.0.1/32")).start));
    try std.testing.expect(isReserved((try parseCidr("0.0.0.0/32")).start));
    try std.testing.expect(isReserved((try parseCidr("192.88.99.1/32")).start));
    try std.testing.expect(!isReserved((try parseCidr("8.8.8.8/32")).start));
    try std.testing.expect(!isReserved((try parseCidr("1.1.1.1/32")).start));
}

test "IpPicker maps indices across user CIDRs minus the reserved floor" {
    const cidrs = [_]Range{
        try parseCidr("8.8.8.0/30"),
        try parseCidr("10.0.0.0/24"),
        try parseCidr("1.1.1.0/31"),
    };
    var picker = try IpPicker.build(std.testing.allocator, &cidrs);
    defer picker.deinit();

    try std.testing.expectEqual(@as(u64, 6), picker.count);
    try std.testing.expectEqual(@as(u32, 0x01010100), picker.at(0));
    try std.testing.expectEqual(@as(u32, 0x01010101), picker.at(1));
    try std.testing.expectEqual(@as(u32, 0x08080800), picker.at(2));
    try std.testing.expectEqual(@as(u32, 0x08080803), picker.at(5));
    var i: u64 = 0;
    while (i < picker.count) : (i += 1) try std.testing.expect(!isReserved(picker.at(i)));
}

test "IpPicker over a fully reserved input is empty" {
    const cidrs = [_]Range{try parseCidr("192.168.0.0/16")};
    var picker = try IpPicker.build(std.testing.allocator, &cidrs);
    defer picker.deinit();
    try std.testing.expectEqual(@as(u64, 0), picker.count);
}

test "Engine is a bijection: every IP:port hit exactly once" {
    const cidrs = [_]Range{ try parseCidr("8.8.8.0/28"), try parseCidr("1.2.3.0/30") };
    const ports = [_]u16{ 80, 443, 22 };
    var eng = try Engine.init(std.testing.allocator, &cidrs, &ports, 0xDEADBEEF);
    defer eng.deinit();

    try std.testing.expectEqual(@as(u64, 60), eng.total);
    var seen = std.AutoHashMap(u64, void).init(std.testing.allocator);
    defer seen.deinit();
    var n: u64 = 0;
    while (eng.next()) |t| {
        try std.testing.expect(!isReserved(t.ip));
        const key = (@as(u64, t.ip) << 16) | t.port;
        try std.testing.expect(!seen.contains(key));
        try seen.put(key, {});
        n += 1;
    }
    try std.testing.expectEqual(@as(u64, 60), n);
    try std.testing.expectEqual(@as(u64, 60), seen.count());
}

test "shards with a shared seed union to the full bijection with no overlap" {
    const cidrs = [_]Range{try parseCidr("8.8.8.0/27")};
    const ports = [_]u16{ 80, 443 };
    const seed: u64 = 0x1234_5678;
    const num_shards: u64 = 4;

    var seen = std.AutoHashMap(u64, void).init(std.testing.allocator);
    defer seen.deinit();
    var emitted: u64 = 0;
    var s: u64 = 0;
    while (s < num_shards) : (s += 1) {
        var eng = try Engine.initShard(std.testing.allocator, &cidrs, &ports, seed, num_shards, s);
        defer eng.deinit();
        while (eng.next()) |t| {
            const key = (@as(u64, t.ip) << 16) | t.port;
            try std.testing.expect(!seen.contains(key));
            try seen.put(key, {});
            emitted += 1;
        }
    }
    try std.testing.expectEqual(@as(u64, 64), emitted);
    try std.testing.expectEqual(@as(u64, 64), seen.count());
}

test "initShard rejects nonsensical shard counts" {
    const cidrs = [_]Range{try parseCidr("8.8.8.0/30")};
    const ports = [_]u16{80};
    try std.testing.expectError(error.InvalidShardCount, Engine.initShard(std.testing.allocator, &cidrs, &ports, 1, 0, 0));
    try std.testing.expectError(error.InvalidShardCount, Engine.initShard(std.testing.allocator, &cidrs, &ports, 1, 100, 0));
    try std.testing.expectError(error.InvalidShardCount, Engine.initShard(std.testing.allocator, &cidrs, &ports, 1, 2, 5));
}

test "parseCidr6 masks host bits, rejects ::/0 and bad prefixes" {
    const c = try parseCidr6("2001:db8:1:2:3:4:5:6/64");
    try std.testing.expectEqual([16]u8{ 0x20, 0x01, 0x0d, 0xb8, 0, 1, 0, 2, 0, 0, 0, 0, 0, 0, 0, 0 }, c.base);
    try std.testing.expectEqual(@as(u8, 64), c.prefix);

    const c120 = try parseCidr6("2001:470:1:2::ab/120");
    try std.testing.expectEqual(@as(u8, 0), c120.base[15]);

    try std.testing.expectError(error.PrefixTooLarge, parseCidr6("::/0"));
    try std.testing.expectError(error.InvalidCidr, parseCidr6("2001:db8::/129"));
    try std.testing.expectError(error.InvalidCidr, parseCidr6("2001:db8::"));
}

test "isReserved6 flags special-use IPv6 blocks and passes global space" {
    try std.testing.expect(isReserved6(try netutil.parseIpv6("::1")));
    try std.testing.expect(isReserved6(try netutil.parseIpv6("::")));
    try std.testing.expect(isReserved6(try netutil.parseIpv6("fe80::1")));
    try std.testing.expect(isReserved6(try netutil.parseIpv6("fc00::1")));
    try std.testing.expect(isReserved6(try netutil.parseIpv6("ff02::1")));
    try std.testing.expect(isReserved6(try netutil.parseIpv6("2001:db8::1")));
    try std.testing.expect(isReserved6(try netutil.parseIpv6("::ffff:c0a8:1")));
    try std.testing.expect(!isReserved6(try netutil.parseIpv6("2001:470:1:2::5")));
    try std.testing.expect(!isReserved6(try netutil.parseIpv6("2606:4700:4700::1111")));
}

test "Engine6 is a bijection over a bounded prefix, every addr:port once, none reserved" {
    const cidr = try parseCidr6("2001:470:1:2::/120");
    const ports = [_]u16{ 80, 443 };
    var eng = try Engine6.init(std.testing.allocator, cidr, &ports, 0xC0FFEE, default_max_hosts6);
    defer eng.deinit();

    try std.testing.expectEqual(@as(u64, 512), eng.total);
    var seen = std.AutoHashMap(u160Key, void).init(std.testing.allocator);
    defer seen.deinit();
    var n: u64 = 0;
    while (eng.next()) |t| {
        try std.testing.expect(!isReserved6(t.addr));
        const key = u160Key{ .addr = t.addr, .port = t.port };
        try std.testing.expect(!seen.contains(key));
        try seen.put(key, {});
        n += 1;
    }
    try std.testing.expectEqual(@as(u64, 512), n);
}

const u160Key = struct { addr: [16]u8, port: u16 };

test "Engine6 rejects prefixes whose host space is too large" {
    const ports = [_]u16{80};
    try std.testing.expectError(error.PrefixTooLarge, Engine6.init(std.testing.allocator, try parseCidr6("2001:470::/64"), &ports, 1, default_max_hosts6));
    try std.testing.expectError(error.PrefixTooLarge, Engine6.init(std.testing.allocator, try parseCidr6("2001:470::/100"), &ports, 1, default_max_hosts6));
    var eng = try Engine6.init(std.testing.allocator, try parseCidr6("2001:470::/112"), &ports, 1, default_max_hosts6);
    eng.deinit();
}

test "Engine6 rejects a host-times-port product that would overflow u64" {
    const ports = [_]u16{ 1, 2, 3, 4 };
    const cidr = try parseCidr6("2001:470::/65");
    try std.testing.expectError(error.PrefixTooLarge, Engine6.init(std.testing.allocator, cidr, &ports, 1, std.math.maxInt(u64)));
}