const std = @import("std");
const mecha = @import("mecha");
const print = std.debug.print;

pub fn solve(part: []u8, buffer: []u8, allocator: std.mem.Allocator) !void {
    var lookup: [256]std.ArrayList(std.meta.Tuple(&.{ i32, i32 })) = undefined;
    for (&lookup) |*element| {
        element.* = std.ArrayList(std.meta.Tuple(&.{ i32, i32 })).init(allocator);
        errdefer element.*.deinit();
    }

    var lines = std.mem.splitScalar(u8, buffer, '\n');
    const width: usize = lines.next().?.len;
    lines.reset();
    const distinct: [][]bool = try allocator.alloc([]bool, width * width);
    defer allocator.free(distinct);
    for (distinct) |*r| {
        r.* = try allocator.alloc(bool, width);
        @memset(r.*, false);
    }

    if (std.mem.eql(u8, part, "1")) {
        try part1(&lookup, distinct, &lines, width);
    } else {
        try part2(&lookup, distinct, &lines, width);
    }

    for (&lookup) |*element| {
        element.*.deinit();
    }

    for (distinct) |r| {
        allocator.free(r);
    }
}

fn part1(lookup: *[256]std.ArrayList(std.meta.Tuple(&.{ i32, i32 })), distinct: [][]bool, lines: *std.mem.SplitIterator(u8, .scalar), width: usize) !void {
    var row: i32 = 0;
    var total_distinct: usize = 0;
    while (lines.next()) |line| : (row += 1) {
        if (line.len == 0) break;

        var col: i32 = 0;
        for (line) |item| {
            if (item == '.' or item == '#') {
                col += 1;
                continue;
            }

            for (lookup[item].items) |antenna| {
                const diffr = row - antenna[0];
                const diffc = col - antenna[1];

                const newr1 = antenna[0] - diffr;
                const newc1 = antenna[1] - diffc;
                if (newr1 >= 0 and newc1 >= 0 and newc1 < width and !distinct[@intCast(newr1)][@intCast(newc1)]) {
                    distinct[@intCast(newr1)][@intCast(newc1)] = true;
                    total_distinct += 1;
                }

                const newr2 = row + diffr;
                const newc2 = col + diffc;
                if (newr2 < width and newc2 >= 0 and newc2 < width and !distinct[@intCast(newr2)][@intCast(newc2)]) {
                    distinct[@intCast(newr2)][@intCast(newc2)] = true;
                    total_distinct += 1;
                }
            }

            try lookup[item].append(.{ row, col });
            col += 1;
        }
    }

    print("{d}\n", .{total_distinct});
}

fn part2(lookup: *[256]std.ArrayList(std.meta.Tuple(&.{ i32, i32 })), distinct: [][]bool, lines: *std.mem.SplitIterator(u8, .scalar), width: usize) !void {
    var row: i32 = 0;
    var total_distinct: usize = 0;
    while (lines.next()) |line| : (row += 1) {
        if (line.len == 0) break;

        var col: i32 = 0;
        for (line) |item| {
            if (item == '.' or item == '#') {
                col += 1;
                continue;
            }

            for (lookup[item].items) |antenna| {
                const diffr = row - antenna[0];
                const diffc = col - antenna[1];

                var newr1 = antenna[0];
                var newc1 = antenna[1];
                while (newr1 >= 0 and newc1 >= 0 and newc1 < width) {
                    if (!distinct[@intCast(newr1)][@intCast(newc1)]) {
                        distinct[@intCast(newr1)][@intCast(newc1)] = true;
                        total_distinct += 1;
                    }
                    newr1 -= diffr;
                    newc1 -= diffc;
                }

                var newr2 = row;
                var newc2 = col;
                while (newr2 < width and newc2 >= 0 and newc2 < width) {
                    if (!distinct[@intCast(newr2)][@intCast(newc2)]) {
                        distinct[@intCast(newr2)][@intCast(newc2)] = true;
                        total_distinct += 1;
                    }
                    newr2 += diffr;
                    newc2 += diffc;
                }
            }

            try lookup[item].append(.{ row, col });
            col += 1;
        }
    }

    print("{d}\n", .{total_distinct});
}