const std = @import("std");

const Pos = struct { x: usize, y: usize };

const Dir = enum { left, right, up, down };

fn step(pos: Pos, dir: Dir) Pos {
    return switch (dir) {
        .left => .{ .x = pos.x -% 1, .y = pos.y },
        .right => .{ .x = pos.x + 1, .y = pos.y },
        .up => .{ .x = pos.x, .y = pos.y -% 1 },
        .down => .{ .x = pos.x, .y = pos.y + 1 },
    };
}

fn invalid(pos: Pos, input: Input) bool {
    return pos.x >= input.width or pos.y >= input.height;
}

const Input = struct {
    plants: []u8,
    width: usize,
    height: usize,
};

t: 
\\ -   
\\ RRRR|
\\|RRRR 
\\   RRR|
\\   R -
\\   -
,

const test_input =
    \\RRRRIICCFF
    \\RRRRIICCCF
    \\VVRRRCCFFF
    \\VVRCCCJFFF
    \\VVVVCJJCFE
    \\VVIVCCJJEE
    \\VVIIICJJEE
    \\MIIIIIJJEE
    \\MIIISIJEEE
    \\MMMISSJEEE
    \\
;

pub fn parse(allocator: std.mem.Allocator, data: []const u8) !Input {
    var it = std.mem.split(u8, data, "\n");
    var plants = std.ArrayList(u8).init(allocator);
    var width: usize = undefined;
    var height: usize = 0;
    while (it.next()) |line| {
        if (line.len == 0) break;
        width = line.len;
        height += 1;
        try plants.appendSlice(line);
    }
    return .{ .plants = plants.items, .width = width, .height = height };
}

test "part1" {
    var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
    defer arena.deinit();

    const output = try part1(arena.allocator(), try parse(arena.allocator(), test_input));
    try std.testing.expectEqual(1930, output);
}

pub fn part1(allocator: std.mem.Allocator, input: Input) !u32 {
    var visited = try allocator.alloc(bool, input.plants.len);
    var stack = std.ArrayList(Pos).init(allocator);
    var cost: u32 = 0;
    for (0..input.height) |start_y| for (0..input.width) |start_x| {
        if (visited[start_y * input.width + start_x]) continue;

        stack.clearRetainingCapacity();
        try stack.append(.{ .x = start_x, .y = start_y });
        visited[start_y * input.width + start_x] = true;
        const plant = input.plants[start_y * input.width + start_x];
        var area: u32 = 1;
        var perimeter: u32 = 0;
        while (stack.popOrNull()) |pos| {
            // std.debug.print("considering {}, {}\n", .{ pos.x, pos.y });
            for ([4]Pos{
                .{ .x = pos.x -% 1, .y = pos.y },
                .{ .x = pos.x + 1, .y = pos.y },
                .{ .x = pos.x, .y = pos.y -% 1 },
                .{ .x = pos.x, .y = pos.y + 1 },
            }) |n| {
                if (n.x >= input.width or n.y >= input.height or input.plants[n.y * input.width + n.x] != plant) {
                    // std.debug.print("1 perimeter from {}, {} towards {}, {}\n", .{ pos.x, pos.y, @as(i64, @bitCast(n.x)), @as(i64, @bitCast(n.y)) });
                    perimeter += 1;
                    continue;
                }
                if (visited[n.y * input.width + n.x]) continue;
                try stack.append(n);
                visited[n.y * input.width + n.x] = true;
                area += 1;
            }
        }
        // std.debug.print("A region of {c} plants with price {} * {}\n", .{ plant, area, perimeter });
        cost += area * perimeter;
    };
    return cost;
}

test "part2" {
    var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
    defer arena.deinit();

    const output = try part2(arena.allocator(), try parse(arena.allocator(), test_input));
    std.debug.print("got {}\n", .{output});
    try std.testing.expectEqual(1206, output);
}

pub fn part2(allocator: std.mem.Allocator, input: Input) !u32 {
    var visited = try allocator.alloc(bool, input.plants.len);
    var stack = std.ArrayList(Pos).init(allocator);
    var cost: u32 = 0;
    for (0..input.height) |start_y| for (0..input.width) |start_x| {
        if (visited[start_y * input.width + start_x]) continue;

        stack.clearRetainingCapacity();
        try stack.append(.{ .x = start_x, .y = start_y });
        visited[start_y * input.width + start_x] = true;
        const plant = input.plants[start_y * input.width + start_x];
        var area: u32 = 1;
        var sides: u32 = 0;
        while (stack.popOrNull()) |pos| {
            // std.debug.print("considering {}, {}\n", .{ pos.x, pos.y });
            for (std.meta.tags(Dir)) |dir| {
                const n = step(pos, dir);
                if (invalid(n, input) or input.plants[n.y * input.width + n.x] != plant) {
                    const side: Dir = switch (dir) {
                        .left => .down,
                        .down => .right,
                        .right => .up,
                        .up => .left,
                    };
                    const s = step(pos, side);
                    const is_leftmost = invalid(s, input) or
                        input.plants[s.y * input.width + s.x] != plant or
                        blk: {
                        const x = step(s, dir);
                        break :blk !invalid(x, input) and input.plants[x.y * input.width + x.x] == plant;
                    };
                    if (is_leftmost) {
                        // std.debug.print("1 side from {}, {} towards {}, {}\n", .{ pos.x, pos.y, @as(i64, @bitCast(n.x)), @as(i64, @bitCast(n.y)) });
                        sides += 1;
                    }
                    continue;
                }
                if (visited[n.y * input.width + n.x]) continue;
                try stack.append(n);
                visited[n.y * input.width + n.x] = true;
                area += 1;
            }
        }
        // std.debug.print("A region of {c} plants with price {} * {}\n", .{ plant, area, sides });
        cost += area * sides;
    };
    return cost;
}