const std = @import("std");

fn cast(n: usize) i32 {
    return @intCast(n);
}

const Input = struct {
    buf: []const u8,
    w: usize,
    h: usize,

    pub fn at(self: Input, x: i32, y: i32) ?u8 {
        if (x < 0 or x >= self.w or y < 0 or y >= self.h) return null;
        return self.buf[@intCast(y * (@as(i32, @intCast(self.w)) + 1) + x)];
    }
};

pub fn parse(_: std.mem.Allocator, data: []const u8) !Input {
    const width = std.mem.indexOf(u8, data, "\n") orelse return error.InvalidInput;
    const height = data.len / (width + 1);
    std.debug.assert((width + 1) * height == data.len);
    return .{ .buf = data, .w = width, .h = 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(),
        \\MMMSXXMASM
        \\MSAMXMSMSA
        \\AMXSXMAAMM
        \\MSAMASMSMX
        \\XMASAMXAMM
        \\XXAMMXXAMA
        \\SMSMSASXSS
        \\SAXAMASAAA
        \\MAMMMXMMMM
        \\MXMXAXMASX
        \\
    ));
    std.debug.print("got {}\n", .{output});
    std.debug.assert(output == 18);
}

pub fn part1(_: std.mem.Allocator, input: Input) !u32 {
    var count: u32 = 0;
    for (0..input.h) |y| for (0..input.w) |x| for ([_]@Vector(2, i32){
        .{ 1, 0 },
        .{ 1, -1 },
        .{ 0, -1 },
        .{ -1, -1 },
        .{ -1, 0 },
        .{ -1, 1 },
        .{ 0, 1 },
        .{ 1, 1 },
    }) |dir| {
        for (0.., "XMAS") |dist, c| {
            if (input.at(cast(x) + dir[0] * cast(dist), cast(y) + cast(dist) * dir[1]) != c) {
                break;
            }
        } else {
            count += 1;
        }
    };

    return count;
}

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

    const output = try part2(arena.allocator(), try parse(arena.allocator(),
        \\MMMSXXMASM
        \\MSAMXMSMSA
        \\AMXSXMAAMM
        \\MSAMASMSMX
        \\XMASAMXAMM
        \\XXAMMXXAMA
        \\SMSMSASXSS
        \\SAXAMASAAA
        \\MAMMMXMMMM
        \\MXMXAXMASX
        \\
    ));
    std.debug.print("got {}\n", .{output});
    std.debug.assert(output == 9);
}

fn isOther(a: ?u8, b: ?u8) bool {
    return switch (a orelse return false) {
        'M' => b == 'S',
        'S' => b == 'M',
        else => false,
    };
}

pub fn part2(_: std.mem.Allocator, input: Input) !u32 {
    var count: u32 = 0;
    for (0..input.h) |y| for (0..input.w) |x| {
        if (input.at(cast(x), cast(y)) != 'A') continue;
        if (!isOther(input.at(cast(x) - 1, cast(y) - 1), input.at(cast(x) + 1, cast(y) + 1))) continue;
        if (!isOther(input.at(cast(x) + 1, cast(y) - 1), input.at(cast(x) - 1, cast(y) + 1))) continue;
        count += 1;
    };

    return count;
}