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 {
    if (std.mem.eql(u8, part, "1")) {
        try part1(buffer, allocator);
    } else {
        try part2(buffer, allocator);
    }
}

fn recursive(mapped: [][]bool, grid: [][]const u8, r: i32, c: i32, region: u8) std.meta.Tuple(&.{ usize, usize }) {
    if (r >= 0 and r < grid.len and c >= 0 and c < grid[@intCast(r)].len and grid[@intCast(r)][@intCast(c)] == region) {
        if (!mapped[@intCast(r)][@intCast(c)]) {
            mapped[@intCast(r)][@intCast(c)] = true;
            const up = recursive(mapped, grid, r - 1, c, region);
            const down = recursive(mapped, grid, r + 1, c, region);
            const left = recursive(mapped, grid, r, c - 1, region);
            const right = recursive(mapped, grid, r, c + 1, region);
            return .{ 1 + up[0] + down[0] + left[0] + right[0], up[1] + down[1] + left[1] + right[1] };
        } else {
            return .{ 0, 0 };
        }
    } else {
        return .{ 0, 1 };
    }
}

fn part1(buffer: []const u8, allocator: std.mem.Allocator) !void {
    var lines = std.mem.splitScalar(u8, buffer, '\n');

    var grid = std.ArrayList([]const u8).init(allocator);
    defer grid.deinit();

    var mapped = std.ArrayList([]bool).init(allocator);
    defer mapped.deinit();

    while (lines.next()) |line| {
        try grid.append(line);

        const mapped_row = try allocator.alloc(bool, line.len);
        @memset(mapped_row, false);
        try mapped.append(mapped_row);
    }

    var output: usize = 0;
    for (grid.items, 0..) |row, ri| {
        for (row, 0..) |item, ci| {
            if (!mapped.items[ri][ci]) {
                const plot = recursive(mapped.items, grid.items, @intCast(ri), @intCast(ci), item);
                output += plot[0] * plot[1];
            }
        }
    }

    for (mapped.items) |mapped_row| {
        allocator.free(mapped_row);
    }

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

//----------------------------------PART2-----------------------------

fn recursive2(mapped: [][]bool, grid: [][]const u8, pr: i32, pc: i32, r: i32, c: i32, region: u8, perimeters: *std.AutoHashMap(std.meta.Tuple(&.{ bool, i32, i32, i32, i32 }), void)) !usize {
    if (r >= 0 and r < grid.len and c >= 0 and c < grid[@intCast(r)].len and grid[@intCast(r)][@intCast(c)] == region) {
        if (!mapped[@intCast(r)][@intCast(c)]) {
            mapped[@intCast(r)][@intCast(c)] = true;
            const up = try recursive2(mapped, grid, r, c, r - 1, c, region, perimeters);
            const down = try recursive2(mapped, grid, r, c, r + 1, c, region, perimeters);
            const left = try recursive2(mapped, grid, r, c, r, c - 1, region, perimeters);
            const right = try recursive2(mapped, grid, r, c, r, c + 1, region, perimeters);
            return 1 + up + down + left + right;
        } else {
            return 0;
        }
    } else {
        const key: std.meta.Tuple(&.{ bool, i32, i32, i32, i32 }) = .{ (c < pc or r < pr), r, c, @intCast(@abs(c - pc)), @intCast(@abs(r - pr)) };
        try perimeters.put(key, {});
        return 0;
    }
}

fn find_edge(b: bool, rr: i32, cc: i32, dr: i32, dc: i32, perimeters: std.AutoHashMap(std.meta.Tuple(&.{ bool, i32, i32, i32, i32 }), void), seen: *std.AutoHashMap(std.meta.Tuple(&.{ bool, i32, i32, i32, i32 }), void)) !void {
    var r = rr + dr;
    var c = cc + dc;
    var key: std.meta.Tuple(&.{ bool, i32, i32, i32, i32 }) = .{ b, r, c, @intCast(@abs(dr)), @intCast(@abs(dc)) };
    while (perimeters.contains(key) and !seen.contains(key)) : (key = .{ b, r, c, @intCast(@abs(dr)), @intCast(@abs(dc)) }) {
        try seen.put(key, {});
        r += dr;
        c += dc;
    }
}

fn part2(buffer: []const u8, allocator: std.mem.Allocator) !void {
    var lines = std.mem.splitScalar(u8, buffer, '\n');

    var grid = std.ArrayList([]const u8).init(allocator);
    defer grid.deinit();

    var mapped = std.ArrayList([]bool).init(allocator);
    defer mapped.deinit();

    while (lines.next()) |line| {
        if (line.len == 0) break;
        try grid.append(line);

        const mapped_row = try allocator.alloc(bool, line.len);
        @memset(mapped_row, false);
        try mapped.append(mapped_row);
    }

    var output: usize = 0;
    for (grid.items, 0..) |row, ri| {
        for (row, 0..) |item, ci| {
            if (!mapped.items[ri][ci]) {
                var perimeters = std.AutoHashMap(std.meta.Tuple(&.{ bool, i32, i32, i32, i32 }), void).init(allocator);
                defer perimeters.deinit();

                const area = try recursive2(mapped.items, grid.items, @intCast(ri), @intCast(ci), @intCast(ri), @intCast(ci), item, &perimeters);
                var seen = std.AutoHashMap(std.meta.Tuple(&.{ bool, i32, i32, i32, i32 }), void).init(allocator);
                defer seen.deinit();
                var iterator = perimeters.keyIterator();
                var edge_count: usize = 0;
                while (iterator.next()) |key| {
                    if (!seen.contains(key.*)) {
                        try find_edge(key[0], key[1], key[2], key[3], key[4], perimeters, &seen);
                        try find_edge(key[0], key[1], key[2], -1 * key[3], -1 * key[4], perimeters, &seen);
                        edge_count += 1;
                    }
                }
                output += area * edge_count;
            }
        }
    }

    for (mapped.items) |mapped_row| {
        allocator.free(mapped_row);
    }

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